fix: make run admission retry-safe

This commit is contained in:
4pmtong 2026-08-05 23:33:37 +08:00
parent d86d057b53
commit ffcf702e18
12 changed files with 411 additions and 56 deletions

View file

@ -13,12 +13,14 @@
# ========= Copyright 2025-2026 @ Eigent.ai All Rights Reserved. =========
import asyncio
import hashlib
import inspect
import json
import logging
import os
from collections.abc import Mapping
from contextlib import suppress
from dataclasses import replace
from dataclasses import dataclass, replace
from pathlib import Path
from dotenv import load_dotenv
@ -88,6 +90,68 @@ SSE_TIMEOUT_SECONDS = 60 * 60
os.environ.setdefault("CAMEL_MODEL_LOG_ENABLED", "true")
@dataclass(frozen=True)
class _PreparedChatRun:
task_lock: TaskLock
run_context: RunContext
attempt_id: str
initial_action: ActionImproveData
def _admission_request_id(
run_id: str,
*,
question: str,
attaches: list[str],
project_context: str | None,
) -> str:
canonical = json.dumps(
{
"question": question,
"attaches": attaches,
"project_context": project_context,
},
ensure_ascii=False,
separators=(",", ":"),
sort_keys=True,
)
digest = hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:24]
return f"initial:{run_id}:{digest}"
async def _classify_persisted_admission(
journal,
*,
run_id: str,
request_id: str,
) -> tuple[str, object | None]:
"""Classify an existing Run as retryable, duplicate, or conflicting."""
run = await asyncio.to_thread(journal.get_run, run_id)
if run is None:
return "new", None
attempts = await asyncio.to_thread(journal.list_run_attempts, run_id)
legacy_request_id = f"initial:{run_id}"
matching = next(
(
attempt
for attempt in attempts
if attempt.resume_request_id in {request_id, legacy_request_id}
),
None,
)
if matching is None:
if not attempts and run.status in {"pending", "running"}:
return "retry", None
return "conflict", None
if run.status in {"pending", "running"} and matching.status in {
"pending",
"running",
}:
return "retry", matching
return "duplicate", matching
def _is_remote_browser_hands(request: Request | None) -> bool:
hands = getattr(getattr(request, "state", None), "hands", None)
if hands is None:
@ -378,7 +442,7 @@ async def _replay_persisted_run(run_id: str):
async def _prepare_chat_run(
data: Chat, request: Request
) -> tuple[TaskLock, RunContext]:
) -> _PreparedChatRun:
"""Perform the one-time compatibility setup for a newly admitted Run."""
# TODO(brain-auth): Phase B should derive canonical user_id from
# request.state.brain_auth, then verify/replace Chat.email before any
@ -439,13 +503,19 @@ async def _prepare_chat_run(
get_default_run_journal().ensure_run,
run_id=run_context.run_id,
project_id=run_context.project_id,
status="pending",
)
await asyncio.to_thread(
attempt = await asyncio.to_thread(
get_default_run_journal().create_run_attempt,
run_context.run_id,
request_id=f"initial:{run_context.run_id}",
request_id=_admission_request_id(
run_context.run_id,
question=data.question,
attaches=data.attaches or [],
project_context=data.project_context,
),
reason="initial_execution",
activate=True,
activate=False,
)
apply_run_env_for_third_party(run_context)
task_lock.run_context = run_context
@ -475,16 +545,22 @@ async def _prepare_chat_run(
# Set the initial current_task_id in task_lock
set_current_task_id(data.project_id, data.task_id)
# Put initial action in queue to start processing
await task_lock.put_queue(
ActionImproveData(
data=ImprovePayload(
question=data.question,
attaches=data.attaches or [],
project_context=data.project_context,
),
new_task_id=data.task_id,
)
request_id = _admission_request_id(
run_context.run_id,
question=data.question,
attaches=data.attaches or [],
project_context=data.project_context,
)
initial_action = ActionImproveData(
data=ImprovePayload(
question=data.question,
attaches=data.attaches or [],
project_context=data.project_context,
),
new_task_id=data.task_id,
request_id=request_id,
run_id=run_context.run_id,
attempt_id=attempt.attempt_id,
)
chat_logger.info(
@ -497,7 +573,12 @@ async def _prepare_chat_run(
"binding_source": frozen_dirs.binding_source,
},
)
return task_lock, run_context
return _PreparedChatRun(
task_lock=task_lock,
run_context=run_context,
attempt_id=attempt.attempt_id,
initial_action=initial_action,
)
async def start_chat_stream(data: Chat, request: Request):
@ -517,25 +598,41 @@ async def start_chat_stream(data: Chat, request: Request):
)
return timeout_stream_wrapper(subscription, run_id=run_id)
persisted_run = await asyncio.to_thread(
journal.get_run,
request_id = _admission_request_id(
run_id,
question=data.question,
attaches=data.attaches or [],
project_context=data.project_context,
)
if persisted_run is not None:
admission, _attempt = await _classify_persisted_admission(
journal,
run_id=run_id,
request_id=request_id,
)
if admission == "conflict":
raise UserException(
code.error,
"This Run id is already bound to a different request.",
)
if admission == "duplicate":
chat_logger.info(
"Replaying persisted Run without implicit restart",
extra={"run_id": run_id, "project_id": data.project_id},
)
return _replay_persisted_run(run_id)
task_lock, run_context = await _prepare_chat_run(data, request)
prepared = await _prepare_chat_run(data, request)
await prepared.task_lock.put_queue(prepared.initial_action)
execution_stream = step_solve(data, request, prepared.task_lock)
subscription = await coordinator.start_with_subscription(
run_id=run_context.run_id,
run_id=prepared.run_context.run_id,
stream_factory=lambda: stream_with_run_context(
step_solve(data, request, task_lock),
lambda: getattr(task_lock, "run_context", run_context),
execution_stream,
lambda: getattr(
prepared.task_lock, "run_context", prepared.run_context
),
),
command_queue=task_lock.queue,
command_queue=prepared.task_lock.queue,
)
return timeout_stream_wrapper(subscription, run_id=run_id)
@ -586,21 +683,42 @@ async def improve(id: str, data: SupplementChat, request: Request):
if data.task_id:
coordinator = get_default_run_coordinator()
async with coordinator.admission_scope(data.task_id):
persisted_run = await asyncio.to_thread(
get_default_run_journal().get_run,
request_id = _admission_request_id(
data.task_id,
question=data.question,
attaches=data.attaches or [],
project_context=data.project_context,
)
if persisted_run is not None:
admission, attempt = await _classify_persisted_admission(
get_default_run_journal(),
run_id=data.task_id,
request_id=request_id,
)
if admission == "conflict":
return Response(status_code=409)
if admission == "duplicate" or (
admission == "retry"
and getattr(attempt, "status", None) == "running"
and await coordinator.get_handle(data.task_id) is not None
):
chat_logger.info(
"Ignored duplicate follow-up Run admission",
extra={"project_id": id, "run_id": data.task_id},
)
return Response(status_code=201)
return await _improve_chat(id, data, request)
return await _improve_chat(
id, data, request, admission_request_id=request_id
)
return await _improve_chat(id, data, request)
async def _improve_chat(id: str, data: SupplementChat, request: Request):
async def _improve_chat(
id: str,
data: SupplementChat,
request: Request,
*,
admission_request_id: str | None = None,
):
chat_logger.info(
"Chat improvement requested",
extra={"task_id": id, "question_length": len(data.question)},
@ -757,13 +875,20 @@ async def _improve_chat(id: str, data: SupplementChat, request: Request):
get_default_run_journal().ensure_run,
run_id=refreshed_context.run_id,
project_id=refreshed_context.project_id,
status="pending",
)
await asyncio.to_thread(
attempt = await asyncio.to_thread(
get_default_run_journal().create_run_attempt,
refreshed_context.run_id,
request_id=f"initial:{refreshed_context.run_id}",
request_id=admission_request_id
or _admission_request_id(
refreshed_context.run_id,
question=data.question,
attaches=data.attaches or [],
project_context=data.project_context,
),
reason="follow_up_execution",
activate=True,
activate=False,
)
await asyncio.to_thread(
get_memory_service().on_run_start,
@ -780,26 +905,22 @@ async def _improve_chat(id: str, data: SupplementChat, request: Request):
prompt_source="improve",
)
if previous_run_id is not None:
await get_default_run_coordinator().rebind_run(
rebound = await get_default_run_coordinator().rebind_run(
previous_run_id,
refreshed_context.run_id,
)
if not rebound:
raise UserException(
code.error,
"The previous Run has no live consumer for this follow-up.",
)
elif data.task_id:
# The client wanted a fresh run but rotation failed upstream. Don't
# touch durable memory; the in-process turn still proceeds so the
# user gets a response, but we leave a breadcrumb for diagnosis.
chat_logger.warning(
"Skipped durable on_run_start: run_context did not rotate to"
" requested task_id",
extra={
"project_id": id,
"requested_task_id": data.task_id,
"current_run_id": (
refreshed_context.run_id
if isinstance(refreshed_context, RunContext)
else None
),
},
raise UserException(
code.error,
"Could not durably prepare the requested follow-up Run.",
)
await task_lock.put_queue(
@ -810,6 +931,11 @@ async def _improve_chat(id: str, data: SupplementChat, request: Request):
project_context=data.project_context,
),
new_task_id=data.task_id,
request_id=admission_request_id,
run_id=(
refreshed_context.run_id if rotation_succeeded else None
),
attempt_id=(attempt.attempt_id if rotation_succeeded else None),
)
)
chat_logger.info(

View file

@ -62,7 +62,7 @@ _TERMINAL_EVENT_TYPES = {
"run.completed",
"run.failed",
"run.cancelled",
"run.timed_out",
"run.deadline_reached",
}

View file

@ -75,9 +75,15 @@ class EventRecorder:
data: dict[str, Any],
event_id: str | None = None,
created_at: float | None = None,
allow_terminal: bool = False,
) -> CommittedRunEvent:
"""Persist one legacy SSE/ChatStep for an already admitted Run."""
if step == "end" and not allow_terminal:
raise ValueError(
"legacy end is reserved for the trusted execution stream"
)
values: dict[str, Any] = {
"event_type": f"legacy.{step}",
"payload": data,

View file

@ -0,0 +1,36 @@
"""Shared in-process activation gate for durable improve commands."""
from __future__ import annotations
import asyncio
import logging
from app.run_journal import get_default_run_journal
from app.service.task import ActionImproveData, TaskLock
async def activate_improve_admission(
task_lock: TaskLock,
item: ActionImproveData,
*,
project_id: str,
logger: logging.Logger,
) -> bool:
"""Activate a pending Attempt once and discard duplicate queue envelopes."""
if not item.request_id:
return True
if item.request_id in task_lock.processed_improve_request_ids:
logger.info(
"Skipping duplicate improve admission",
extra={"project_id": project_id, "request_id": item.request_id},
)
return False
if item.attempt_id and item.run_id:
await asyncio.to_thread(
get_default_run_journal().activate_run_attempt,
item.attempt_id,
expected_run_id=item.run_id,
)
task_lock.processed_improve_request_ids.add(item.request_id)
return True

View file

@ -65,6 +65,7 @@ from app.memory import (
)
from app.model.chat import Chat, NewAgent, Status, TaskContent, sse_json
from app.model.subscription_runtime import is_subscription_auth
from app.run_runtime.admission import activate_improve_admission
from app.service.single_agent_service import single_agent_solve
from app.service.task import (
Action,
@ -96,6 +97,20 @@ SUMMARY_TASK_NAME_MAX_LENGTH = 80
SUMMARY_TASK_SUMMARY_MAX_LENGTH = 240
async def _activate_improve_admission(
task_lock: TaskLock,
item: ActionImproveData,
*,
project_id: str,
) -> bool:
return await activate_improve_admission(
task_lock,
item,
project_id=project_id,
logger=logger,
)
def _truncate_summary_part(value: str, max_length: int) -> str:
text = " ".join((value or "").replace("|", " ").split())
if len(text) <= max_length:
@ -548,6 +563,14 @@ async def step_solve(options: Chat, request: Request, task_lock: TaskLock):
# Continue waiting instead of breaking on queue error
continue
if isinstance(item, ActionImproveData):
if not await _activate_improve_admission(
task_lock,
item,
project_id=options.project_id,
):
continue
try:
if item.action == Action.improve or start_event_loop:
logger.info("=" * 80)

View file

@ -29,6 +29,7 @@ from app.memory import (
)
from app.model.chat import Chat, sse_json
from app.model.enums import Status
from app.run_runtime.admission import activate_improve_admission
from app.service.task import (
Action,
ActionData,
@ -304,6 +305,13 @@ async def single_agent_solve(
if item.action == Action.improve:
assert isinstance(item, ActionImproveData)
if not await activate_improve_admission(
task_lock,
item,
project_id=options.project_id,
logger=logger,
):
continue
if item.new_task_id:
current_task_id = item.new_task_id
set_current_task_id(

View file

@ -88,6 +88,9 @@ class ActionImproveData(BaseModel):
action: Literal[Action.improve] = Action.improve
data: ImprovePayload
new_task_id: str | None = None
request_id: str | None = None
run_id: str | None = None
attempt_id: str | None = None
class ActionStartData(BaseModel):
@ -419,6 +422,8 @@ class TaskLock:
"""Latest local history persistence error for diagnostics."""
_memory_finalized_runs: set[str]
"""Run ids whose durable memory lifecycle has already been finalized."""
processed_improve_request_ids: set[str]
"""In-process dedupe for durable admission retries that enqueue twice."""
def __init__(
self, id: str, queue: asyncio.Queue, human_input: dict
@ -453,6 +458,7 @@ class TaskLock:
self.base_snapshot_id = None
self.new_folder_path = None
self.memory_service = None
self.processed_improve_request_ids = set()
self.local_history_degraded = False
self.local_history_last_error = None
self._memory_finalized_runs = set()

View file

@ -202,6 +202,7 @@ async def _record_local_step(args, value) -> None:
run_id=run_id,
step=data["step"],
data=data["data"],
allow_terminal=data["step"] == "end",
)

View file

@ -24,6 +24,10 @@ from fastapi.testclient import TestClient
from pydantic import ValidationError
from app.controller.chat_controller import (
_PreparedChatRun,
_admission_request_id,
_classify_persisted_admission,
_classify_persisted_admission,
human_reply,
improve,
install_mcp,
@ -36,6 +40,8 @@ from app.controller.chat_controller import (
from app.exception.exception import UserException
from app.model.chat import Chat, HumanReply, McpServers, Status, SupplementChat
from app.run_context import RunContext
from app.run_journal import SQLiteRunJournal
from app.run_journal import SQLiteRunJournal
from app.run_runtime import RunCoordinator
@ -43,6 +49,14 @@ from app.run_runtime import RunCoordinator
def controller_run_journal():
journal = MagicMock()
journal.get_run.return_value = None
journal.create_run_attempt.return_value = SimpleNamespace(
attempt_id="attempt-1", status="pending"
)
journal.list_run_attempts.return_value = []
journal.create_run_attempt.return_value = SimpleNamespace(
attempt_id="attempt-1",
status="pending",
)
with patch(
"app.controller.chat_controller.get_default_run_journal",
return_value=journal,
@ -97,6 +111,7 @@ class TestChatController:
controller_run_journal.ensure_run.assert_called_once_with(
run_id=chat_data.run_id or chat_data.task_id,
project_id=chat_data.project_id,
status="pending",
)
@pytest.mark.asyncio
@ -115,7 +130,14 @@ class TestChatController:
await release.wait()
yield "data: once\n\n"
prepare = AsyncMock(return_value=(mock_task_lock, run_context))
prepare = AsyncMock(
return_value=_PreparedChatRun(
task_lock=mock_task_lock,
run_context=run_context,
attempt_id="attempt-1",
initial_action=MagicMock(),
)
)
with (
patch(
"app.controller.chat_controller.get_default_run_coordinator",
@ -153,6 +175,47 @@ class TestChatController:
await retry_stream.aclose()
await coordinator.close()
@pytest.mark.asyncio
async def test_pending_partial_admission_is_retryable_but_reuse_conflicts(
self, tmp_path
):
run_id = "run-partial"
request_id = _admission_request_id(
run_id,
question="original",
attaches=[],
project_context=None,
)
with SQLiteRunJournal(tmp_path / "journal.sqlite3") as journal:
journal.ensure_run(
run_id=run_id, project_id="project-1", status="pending"
)
journal.create_run_attempt(
run_id,
request_id=request_id,
reason="initial_execution",
activate=False,
)
retry, _attempt = await _classify_persisted_admission(
journal,
run_id=run_id,
request_id=request_id,
)
conflict, _attempt = await _classify_persisted_admission(
journal,
run_id=run_id,
request_id=_admission_request_id(
run_id,
question="different",
attaches=[],
project_context=None,
),
)
assert retry == "retry"
assert conflict == "conflict"
@pytest.mark.asyncio
async def test_persisted_run_replays_without_implicit_restart(
self,
@ -165,8 +228,19 @@ class TestChatController:
coordinator = RunCoordinator()
controller_run_journal.get_run.return_value = SimpleNamespace(
run_id=run_id,
status="running",
status="completed",
)
controller_run_journal.list_run_attempts.return_value = [
SimpleNamespace(
resume_request_id=_admission_request_id(
run_id,
question=chat_data.question,
attaches=chat_data.attaches,
project_context=chat_data.project_context,
),
status="completed",
)
]
controller_run_journal.list_events.return_value = [
SimpleNamespace(
legacy_step="end",
@ -523,6 +597,7 @@ class TestChatController:
controller_run_journal.ensure_run.assert_called_once_with(
run_id="run-new",
project_id="project-1",
status="pending",
)
assert await coordinator.get_handle("run-old") is None
assert await coordinator.get_handle("run-new") is subscription.handle
@ -540,8 +615,19 @@ class TestChatController:
):
data = SupplementChat(question="duplicate", task_id="run-existing")
controller_run_journal.get_run.return_value = SimpleNamespace(
run_id="run-existing"
run_id="run-existing", status="completed"
)
controller_run_journal.list_run_attempts.return_value = [
SimpleNamespace(
resume_request_id=_admission_request_id(
"run-existing",
question=data.question,
attaches=data.attaches,
project_context=data.project_context,
),
status="completed",
)
]
with patch(
"app.controller.chat_controller._improve_chat",
@ -634,14 +720,9 @@ class TestChatController:
new=AsyncMock(return_value=True),
),
):
response = await improve(
"project_x", supplement_data, mock_request
)
with pytest.raises(UserException, match="durably prepare"):
await improve("project_x", supplement_data, mock_request)
assert isinstance(response, Response)
# The improve request itself still succeeds -- chat must not break
# because durable memory is unhappy.
assert response.status_code == 201
# Critical assertion: on_run_start was NOT called against the stale
# context. The R26 fix only checked data.task_id; R27 strengthens
# it to compare refreshed_context.run_id == data.task_id.

View file

@ -21,6 +21,7 @@ from unittest.mock import MagicMock, patch
import pytest
from app.controller.run_controller import (
_is_terminal,
get_run,
get_run_events,
stream_run_events,
@ -56,6 +57,21 @@ def _event(sequence: int, step: str) -> CommittedRunEvent:
)
def test_deadline_reached_is_a_terminal_stream_event():
event = CommittedRunEvent(
event_id="deadline",
run_id="run-1",
sequence=1,
event_type="run.deadline_reached",
payload={},
legacy_step=None,
created_at=1.0,
run_version=1,
)
assert _is_terminal(event) is True
def _decode_sse(value: str) -> tuple[int | None, str, dict]:
event_id = None
event_name = ""

View file

@ -457,3 +457,26 @@ async def test_event_recorder_rejects_cross_project_attribution(journal):
data={"message": "original"},
event_id="event-1",
)
@pytest.mark.asyncio
async def test_legacy_end_requires_trusted_execution_stream(journal):
journal.ensure_run(run_id="run-1", project_id="project-1")
recorder = EventRecorder(journal)
with pytest.raises(ValueError, match="trusted execution stream"):
await recorder.record_legacy_step(
project_id="project-1",
run_id="run-1",
step="end",
data={},
)
await recorder.record_legacy_step(
project_id="project-1",
run_id="run-1",
step="end",
data={},
allow_terminal=True,
)
assert journal.get_run("run-1").status == "completed"

View file

@ -22,6 +22,7 @@ from camel.tasks.task import TaskState
from app.model.chat import AgentModelConfig, Chat, NewAgent
from app.service.chat_service import (
_activate_improve_admission,
_extract_stream_chunk_content,
_render_subtask_report,
_trim_in_process_history,
@ -74,6 +75,34 @@ class _AgentStepResponse:
self.msgs = [MagicMock(content=content)]
@pytest.mark.asyncio
async def test_improve_admission_activates_once_and_deduplicates_retry():
task_lock = MagicMock()
task_lock.processed_improve_request_ids = set()
journal = MagicMock()
item = ActionImproveData(
data=ImprovePayload(question="hello"),
request_id="request-1",
run_id="run-1",
attempt_id="attempt-1",
)
with patch(
"app.run_runtime.admission.get_default_run_journal",
return_value=journal,
):
assert await _activate_improve_admission(
task_lock, item, project_id="project-1"
)
assert not await _activate_improve_admission(
task_lock, item, project_id="project-1"
)
journal.activate_run_attempt.assert_called_once_with(
"attempt-1", expected_run_id="run-1"
)
@pytest.mark.unit
class TestExtractStreamChunkContent:
def test_extracts_single_message_content(self):