diff --git a/astrbot/core/agent/context/compressor.py b/astrbot/core/agent/context/compressor.py index 759604dd93..8f60e29723 100644 --- a/astrbot/core/agent/context/compressor.py +++ b/astrbot/core/agent/context/compressor.py @@ -1,10 +1,14 @@ +import asyncio +from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Protocol, runtime_checkable from ...provider.modalities import ( log_context_sanitize_stats, sanitize_contexts_by_modalities, ) -from ..message import Message +from ..event_stream import RequestEventRecorder, request_recorder_kwargs +from ..message import Message, dump_messages_with_checkpoints +from ..response import AgentResponse from .token_counter import EstimateTokenCounter, TokenCounter if TYPE_CHECKING: @@ -130,6 +134,8 @@ def __init__( instruction_text: str | None = None, compression_threshold: float = 0.82, token_counter: TokenCounter | None = None, + request_event_emitter: Callable[[AgentResponse], Awaitable[None]] | None = None, + turn_id: str | None = None, ) -> None: """Initialize the LLM summary compressor. @@ -139,8 +145,13 @@ def __init__( exact context. Clamped to 0-0.3. instruction_text: Custom instruction for summary generation. compression_threshold: The compression trigger threshold (default: 0.82). + token_counter: Optional context token estimator. + request_event_emitter: Optional acknowledged runner event sink. + turn_id: Runner turn identity attached to request events. """ self.provider = provider + self.request_event_emitter = request_event_emitter + self.turn_id = turn_id self.keep_recent_ratio = min(max(float(keep_recent_ratio), 0.0), 0.3) self.compression_threshold = compression_threshold self.token_counter = token_counter or EstimateTokenCounter() @@ -247,6 +258,16 @@ async def __call__(self, messages: list[Message]) -> list[Message]: if not any(msg.role != "system" for msg in summary_contexts): return messages + summary_contexts = [ + Message.model_validate(item) + for item in dump_messages_with_checkpoints( + [message for message in summary_contexts if not message._no_save] + ) + if item.get("role") != "_checkpoint" + ] + if not summary_contexts: + return messages + if summary_contexts[-1].role != "assistant": summary_contexts.append( Message( @@ -273,9 +294,37 @@ async def __call__(self, messages: list[Message]) -> list[Message]: # Generate summary try: - response = await self.provider.text_chat( - contexts=sanitized_summary_contexts, - ) + recorder = None + if self.request_event_emitter: + recorder = RequestEventRecorder( + self.request_event_emitter, + { + "turn_id": self.turn_id, + "purpose": "compaction", + "provider_id": self.provider.provider_config.get("id", ""), + "model": self.provider.get_model(), + }, + ) + await recorder.begin() + try: + response = await self.provider.text_chat( + contexts=sanitized_summary_contexts, + **request_recorder_kwargs(self.provider.text_chat, recorder), + ) + if recorder: + await recorder.finish( + "failed" if response.role == "err" else "completed", + usage=response.usage, + ) + except BaseException as exc: + if recorder: + await recorder.finish( + "cancelled" + if isinstance(exc, asyncio.CancelledError) + else "failed", + error_code=type(exc).__name__, + ) + raise summary_content = (response.completion_text or "").strip() except Exception as e: logger.error(f"Failed to generate summary: {e}") diff --git a/astrbot/core/agent/context/config.py b/astrbot/core/agent/context/config.py index aa216d9a25..0d5def1169 100644 --- a/astrbot/core/agent/context/config.py +++ b/astrbot/core/agent/context/config.py @@ -1,6 +1,8 @@ +from collections.abc import Awaitable, Callable from dataclasses import dataclass from typing import TYPE_CHECKING +from ..response import AgentResponse from .compressor import ContextCompressor from .token_counter import TokenCounter @@ -33,3 +35,7 @@ class ContextConfig: """Custom token counting method. If None, the default method is used.""" custom_compressor: ContextCompressor | None = None """Custom context compression method. If None, the default method is used.""" + request_event_emitter: Callable[[AgentResponse], Awaitable[None]] | None = None + """Optional runtime event sink for summary model requests.""" + turn_id: str | None = None + """Runner turn identity for summary request events.""" diff --git a/astrbot/core/agent/context/manager.py b/astrbot/core/agent/context/manager.py index 1a11ebff96..5df6d6d3c0 100644 --- a/astrbot/core/agent/context/manager.py +++ b/astrbot/core/agent/context/manager.py @@ -36,6 +36,8 @@ def __init__( keep_recent_ratio=config.llm_compress_keep_recent_ratio, instruction_text=config.llm_compress_instruction, token_counter=self.token_counter, + request_event_emitter=config.request_event_emitter, + turn_id=config.turn_id, ) else: self.compressor = TruncateByTurnsCompressor( diff --git a/astrbot/core/agent/conversation_events.py b/astrbot/core/agent/conversation_events.py new file mode 100644 index 0000000000..837d2c2fc2 --- /dev/null +++ b/astrbot/core/agent/conversation_events.py @@ -0,0 +1,498 @@ +"""Host-side event collection and branch-bound conversation commits.""" + +import asyncio +import hashlib +import json +import uuid +from collections import Counter, defaultdict, deque +from contextvars import ContextVar +from copy import deepcopy + +from sqlalchemy.exc import SQLAlchemyError + +from astrbot.core.agent.response import AgentResponse +from astrbot.core.db.conversation import ( + CONTEXT_TYPES, + ConversationSnapshot, + persistent_messages, +) +from astrbot.core.sentinels import NOT_GIVEN + +active_plugin_id: ContextVar[str | None] = ContextVar( + "active_conversation_plugin", default=None +) + +active_conversation_writer: ContextVar["ConversationEventWriter | None"] = ContextVar( + "active_conversation_writer", default=None +) + + +class ConversationPersistenceError(RuntimeError): + """A journal write failed; provider or tool retries cannot repair it.""" + + +class ConversationEventWriter: + """Collect runner and plugin events against one conversation revision. + + A runner may retain its own journal. Reuse event IDs when retrying a write, + and never replay tool side effects to resolve a storage conflict. + """ + + def __init__(self, store, snapshot: ConversationSnapshot): + self.store = store + self.cid = snapshot.conversation.conversation_id + self.umo = snapshot.conversation.umo + self.head_seq = snapshot.conversation.head_seq + self.leaf_event_id = snapshot.conversation.leaf_event_id + self.turn_id: str | None = None + self.closed = False + self.turn_result: AgentResponse | None = None + self._lock = asyncio.Lock() + self._pending: list[dict] = [] + self._entries = deepcopy(snapshot.entries) + self._staged_leaf = self.leaf_event_id + self._excluded_counts: Counter[str] = Counter() + self.runtime_context = None + self.request = None + + async def consume(self, response: AgentResponse) -> bool: + """Collect a runner event before allowing execution to resume. + + Args: + response: A value yielded by the runner, with optional stable identity. + + Returns: + Whether the response belongs to persistence rather than display. + """ + if response.type == "context.updated": + self.stage_history( + response.data["messages"], + reason=response.data.get("reason", "legacy_replace"), + ) + elif response.type == "turn.started": + await self.start_turn( + {"umo": self.umo, **response.data.get("trigger", {})}, + event_id=response.event_id, + ) + elif response.type == "turn.finished": + # The host settles the turn after its history-save hook has completed. + self.turn_result = response + elif response.type in { + "request.started", + "request.finished", + "tool.started", + "tool.finished", + "message.appended", + "context.rebased", + }: + payload = dict(response.data) + if response.type in { + "request.started", + "tool.started", + "message.appended", + "context.rebased", + }: + payload.setdefault("turn_id", self.turn_id) + if response.type == "request.started": + payload.setdefault("context_leaf_event_id", self.leaf_event_id) + await self.append( + response.type, payload, event_id=response.event_id, sync_runtime=False + ) + else: + return False + return True + + async def append( + self, + event_type: str, + payload: dict, + *, + event_id=None, + parent_event_id=NOT_GIVEN, + sync_runtime=True, + ): + """Commit one common protocol event without exposing database internals. + + Args: + event_type: Core or namespaced plugin event type. + payload: JSON-serializable event data. + event_id: Stable identity reused for delivery retries. + parent_event_id: Omit to append to the bound branch; None means root. + sync_runtime: Apply plugin writes to working context; runner events already did so. + + Returns: + The acknowledged immutable event. + """ + if self.closed: + raise ValueError("The turn writer is closed") + draft = { + "type": event_type, + "payload": payload, + "event_id": event_id or str(uuid.uuid4()), + } + if parent_event_id is not NOT_GIVEN: + draft["parent_event_id"] = parent_event_id + async with self._lock: + context_change = event_type in CONTEXT_TYPES + drafts = [*self._pending, draft] if context_change else [draft] + previous_staged_leaf = self._staged_leaf + try: + committed = await self.store.append( + self.cid, + drafts, + expected_head=self.head_seq, + expected_leaf=self.leaf_event_id, + ) + except (SQLAlchemyError, ValueError) as exc: + raise ConversationPersistenceError( + "Conversation event was not acknowledged" + ) from exc + acknowledged = next(e for e in committed if e.event_id == draft["event_id"]) + if all(e.seq <= self.head_seq for e in committed): + return acknowledged + self.head_seq = max(self.head_seq, *(e.seq for e in committed)) + for event in committed: + if event.type in CONTEXT_TYPES: + self.leaf_event_id = event.event_id + if context_change: + # Read this immutable tip, even if another writer has since selected a branch. + async with self.store.db.get_db() as session: + from sqlmodel import select + + from astrbot.core.db.po import ConversationV3 + + conv = ( + await session.execute( + select(ConversationV3).where( + ConversationV3.conversation_id == self.cid + ) + ) + ).scalar_one() + snapshot = await self.store.project( + session, conv, self.leaf_event_id + ) + self._entries = snapshot.entries + self._pending.clear() + self._staged_leaf = self.leaf_event_id + if sync_runtime and self.runtime_context is not None: + from astrbot.core.agent.message import Message + + if event_type == "message.appended" and ( + parent_event_id is NOT_GIVEN + or parent_event_id == previous_staged_leaf + ): + message = Message.model_validate(deepcopy(payload["message"])) + message._no_save = not payload.get("include_in_context", True) + self.runtime_context.messages.append(message) + if message._no_save: + excluded = deepcopy(payload["message"]) + excluded["_no_save"] = True + fingerprint = hashlib.sha256( + json.dumps(excluded, sort_keys=True).encode() + ).hexdigest() + self._excluded_counts[fingerprint] += 1 + else: + leading_system = self.runtime_context.messages[:1] + if not leading_system or leading_system[0].role != "system": + leading_system = [] + self.runtime_context.messages = leading_system + [ + Message.model_validate(m) for m in snapshot.messages + ] + elif sync_runtime and self.request is not None: + self.request.contexts = snapshot.messages + return acknowledged + + async def append_message( + self, + message: dict, + *, + include_in_context=True, + event_id=None, + parent_event_id=NOT_GIVEN, + ): + """Persist a message and synchronize the runner's working context. + + Args: + message: AstrBot message dictionary. + include_in_context: Whether future projections include the message. + event_id: Stable delivery ID. + parent_event_id: Optional explicit branch parent. + + Returns: + Committed message event. + """ + return await self.append( + "message.appended", + { + "message": message, + "turn_id": self.turn_id, + "include_in_context": include_in_context + and not message.get("_no_save", False), + }, + event_id=event_id, + parent_event_id=parent_event_id, + ) + + async def start_turn(self, trigger: dict, *, event_id=None) -> str: + """Start this runner's turn once. + + Args: + trigger: Origin metadata, without request content. + event_id: Optional host-assigned turn identity. + + Returns: + Stable turn identity. + """ + if self.turn_id is not None: + return self.turn_id + event = await self.append( + "turn.started", + { + "trigger": trigger, + "base_leaf_event_id": self.leaf_event_id, + }, + event_id=event_id, + ) + self.turn_id = event.event_id + return self.turn_id + + async def finish_turn(self, status: str) -> None: + """Record settlement without claiming recovery of external side effects. + + Args: + status: completed, failed, or cancelled. + """ + if self.turn_id and not self.closed: + payload = dict(self.turn_result.data) if self.turn_result else {} + payload.update(turn_id=self.turn_id, status=status) + await self.append( + "turn.finished", + payload, + event_id=self.turn_result.event_id if self.turn_result else None, + ) + self.closed = True + self._pending.clear() + + def stage_history(self, history: list[dict], *, reason="legacy_replace") -> None: + """Stage legacy mutations until the existing history-save boundary. + + Args: + history: Complete working history, including projection exclusions. + reason: Reason for a replacement, such as compaction. + """ + excluded_counts = Counter() + for message in history: + effective = persistent_messages([message]) + if message.get("role") == "_checkpoint" or effective == [message]: + continue + fingerprint = hashlib.sha256( + json.dumps(message, sort_keys=True).encode() + ).hexdigest() + excluded_counts[fingerprint] += 1 + if excluded_counts[fingerprint] <= self._excluded_counts[fingerprint]: + continue + identity = str(uuid.uuid4()) + self._pending.append( + { + "event_id": identity, + "type": "message.appended", + "parent_event_id": self._staged_leaf, + "payload": { + "turn_id": self.turn_id, + "message": deepcopy(message), + "include_in_context": False, + }, + } + ) + self._staged_leaf = identity + self._excluded_counts |= excluded_counts + history = persistent_messages(history) + before = [e["message"] for e in self._entries] + if before == history: + return + if history[: len(before)] == before: + for message in history[len(before) :]: + identity = str(uuid.uuid4()) + self._pending.append( + { + "event_id": identity, + "type": "message.appended", + "parent_event_id": self._staged_leaf, + "payload": { + "turn_id": self.turn_id, + "message": deepcopy(message), + }, + } + ) + self._entries.append({"id": identity, "message": deepcopy(message)}) + self._staged_leaf = identity + else: + identity = str(uuid.uuid4()) + retained_ids = defaultdict(deque) + for entry in self._entries: + retained_ids[json.dumps(entry["message"], sort_keys=True)].append( + entry["id"] + ) + self._entries = [] + for message in history: + matches = retained_ids[json.dumps(message, sort_keys=True)] + self._entries.append( + { + "id": matches.popleft() if matches else str(uuid.uuid4()), + "message": deepcopy(message), + } + ) + self._pending.append( + { + "event_id": identity, + "type": "context.rebased", + "parent_event_id": self._staged_leaf, + "payload": { + "reason": "reset" if not history else reason, + "turn_id": self.turn_id, + "messages": deepcopy(self._entries), + }, + } + ) + self._staged_leaf = identity + + async def save_history(self, history: list[dict], *, token_usage=None) -> None: + """Commit staged context at the old save boundary with revision checks. + + Args: + history: Final legacy working history. + token_usage: Optional legacy current-context estimate. + """ + async with self._lock: + self.stage_history(history) + drafts = list(self._pending) + if token_usage is not None: + drafts.append( + { + "type": "conversation.updated", + "payload": { + "changes": {}, + "token_usage": token_usage, + }, + } + ) + committed = await self.store.append( + self.cid, + drafts, + expected_head=self.head_seq, + expected_leaf=self.leaf_event_id, + ) + for event in committed: + self.head_seq = event.seq + if event.type in CONTEXT_TYPES: + self.leaf_event_id = event.event_id + self._staged_leaf = self.leaf_event_id + self._pending.clear() + if self.runtime_context is not None: + from astrbot.core.agent.message import bind_checkpoint_messages + + system = self.runtime_context.messages[:1] + if not system or system[0].role != "system": + system = [] + self.runtime_context.messages = system + bind_checkpoint_messages( + deepcopy(history) + ) + elif self.request is not None: + self.request.contexts = deepcopy(history) + + def plugin(self, plugin_id: str) -> "PluginConversationEvents": + """Create a namespace-bound plugin journal facade. + + Args: + plugin_id: Identity resolved by the plugin loader, not payload data. + + Returns: + A plugin-specific append/read facade. + """ + if not plugin_id or any( + c not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-." + for c in plugin_id + ): + raise ValueError("Invalid plugin identity") + return PluginConversationEvents(self, plugin_id) + + +class PluginConversationEvents: + """Plugin context writes and a namespace-bound private event journal.""" + + def __init__(self, writer: ConversationEventWriter, plugin_id: str): + self.writer = writer + self.prefix = f"plugin.{plugin_id}." + + async def append_message( + self, + message: dict, + *, + include_in_context=True, + event_id=None, + parent_event_id=NOT_GIVEN, + ): + """Append a core message and synchronize the current working context. + + Args: + message: AstrBot message dictionary. + include_in_context: Whether future projections include the message. + event_id: Stable delivery identity for retries. + parent_event_id: Omit for the current branch; None starts a root. + + Returns: + Committed message event. + """ + return await self.writer.append_message( + message, + include_in_context=include_in_context, + event_id=event_id, + parent_event_id=parent_event_id, + ) + + async def append(self, name: str, payload: dict, *, event_id=None): + """Commit plugin-private JSON data. + + Args: + name: Plugin-local event name. + payload: Durable plugin data; do not copy temporary prompt content. + event_id: Stable delivery identity for retries. + + Returns: + Committed event. + """ + if not name or "." in name: + raise ValueError("Use a local event name without dots") + return await self.writer.append(self.prefix + name, payload, event_id=event_id) + + async def list(self, name: str, *, after_seq=0, limit=100): + """Read a bounded page in this conversation and namespace. + + Args: + name: Plugin-local event name. + after_seq: Exclusive cursor. + limit: Maximum event count. + + Returns: + Stored private events. + """ + return await self.writer.store.events( + self.writer.cid, + event_type=self.prefix + name, + after_seq=after_seq, + limit=limit, + ) + + async def latest(self, name: str): + """Read the most recent private event without scanning history. + + Args: + name: Plugin-local event name. + + Returns: + Last event, or None. + """ + events = await self.writer.store.events( + self.writer.cid, event_type=self.prefix + name, descending=True, limit=1 + ) + return events[0] if events else None diff --git a/astrbot/core/agent/event_stream.py b/astrbot/core/agent/event_stream.py new file mode 100644 index 0000000000..a467d36622 --- /dev/null +++ b/astrbot/core/agent/event_stream.py @@ -0,0 +1,153 @@ +"""Storage-independent runner events and acknowledged delivery.""" + +import asyncio +import inspect +import uuid +from collections.abc import AsyncGenerator, Awaitable, Callable + +from .response import AgentResponse + + +class AgentEventStream: + """Expose nested async operations through one backpressured response stream. + + The producer resumes only after the consumer requests the next response. + Closing the stream cancels pending work before acknowledging another event. + """ + + def __init__(self): + self.active = False + self._queue: asyncio.Queue[tuple[AgentResponse, asyncio.Future[None]]] = ( + asyncio.Queue(maxsize=1) + ) + + async def emit(self, response: AgentResponse) -> None: + """Wait until the host has handled a response. + + Args: + response: Runtime or durable event to deliver to the host. + """ + if not self.active: + raise asyncio.CancelledError + acknowledged = asyncio.get_running_loop().create_future() + await self._queue.put((response, acknowledged)) + await acknowledged + + async def run(self, source: AsyncGenerator[AgentResponse, None]): + """Drive one step and join its producer even when consumption stops early. + + Args: + source: The runner's step implementation. + + Yields: + Responses in execution order, one at a time. + """ + if self.active: + raise RuntimeError("A runner step is already being consumed") + self.active = True + + async def produce(): + try: + async for response in source: + await self.emit(response) + finally: + await source.aclose() + + producer = asyncio.create_task(produce()) + receiving = None + try: + while True: + receiving = asyncio.create_task(self._queue.get()) + done, _ = await asyncio.wait( + {producer, receiving}, return_when=asyncio.FIRST_COMPLETED + ) + if receiving in done: + response, acknowledged = receiving.result() + yield response + if not acknowledged.done(): + acknowledged.set_result(None) + else: + await producer + break + finally: + self.active = False + if receiving is not None: + receiving.cancel() + producer.cancel() + await asyncio.gather( + producer, + *([receiving] if receiving is not None else []), + return_exceptions=True, + ) + while not self._queue.empty(): + _, acknowledged = self._queue.get_nowait() + acknowledged.cancel() + + +def request_recorder_kwargs(call: Callable, recorder) -> dict: + """Pass instrumentation only to adapters that explicitly accept it. + + Args: + call: Provider method, which may be implemented by an older plugin. + recorder: Optional request attempt recorder. + + Returns: + Local instrumentation arguments, never arbitrary model request kwargs. + """ + if ( + recorder is not None + and "request_event_recorder" in inspect.signature(call).parameters + ): + return {"request_event_recorder": recorder} + return {} + + +class RequestEventRecorder: + """Pair provider attempts while attaching final usage after response decoding.""" + + def __init__( + self, emit: Callable[[AgentResponse], Awaitable[object]], payload: dict + ): + """Attach an explicit runtime sink to one provider invocation. + + Args: + emit: Awaited delivery of events to the runner's consumer. + payload: Request metadata, excluding the model request body. + """ + self.emit = emit + self.payload = payload + self.current_id = None + self.network_attempt_seen = False + + async def begin(self): + """Start an attempt, retaining its ID until settlement.""" + self.current_id = str(uuid.uuid4()) + await self.emit(AgentResponse("request.started", self.payload, self.current_id)) + + async def before_network_attempt(self): + """Open another attempt when the shared provider retry loop retries.""" + if self.network_attempt_seen and self.current_id is None: + await self.begin() + self.network_attempt_seen = True + + async def finish(self, status, *, usage=None, error_code=None): + """Settle one attempt once without storing its request body. + + Args: + status: Terminal request status. + usage: Optional normalized token usage from the decoded response. + error_code: Optional exception class or stable failure code. + """ + if self.current_id is None: + return + payload = {"request_id": self.current_id, "status": status} + if usage is not None: + payload["usage"] = { + "input_tokens": usage.input, + "cached_input_tokens": usage.input_cached, + "output_tokens": usage.output, + } + if error_code: + payload["error"] = {"code": error_code} + await self.emit(AgentResponse("request.finished", payload, str(uuid.uuid4()))) + self.current_id = None diff --git a/astrbot/core/agent/message.py b/astrbot/core/agent/message.py index 4292f4c04e..68f20e575d 100644 --- a/astrbot/core/agent/message.py +++ b/astrbot/core/agent/message.py @@ -342,16 +342,30 @@ def bind_checkpoint_messages(history: list[dict]) -> list[Message]: return messages -def dump_messages_with_checkpoints(messages: list[Message]) -> list[dict]: - """Dump runtime messages and reinsert bound checkpoint segments.""" +def dump_messages_with_checkpoints( + messages: list[Message], *, include_temporary: bool = False +) -> list[dict]: + """Dump runtime messages and reinsert bound checkpoint segments. + + Args: + messages: Runtime working messages. + include_temporary: Preserve exclusion markers for the event journal. + + Returns: + Serialized messages with optional temporary-content markers. + """ dumped: list[dict] = [] for message in messages: message_data = message.model_dump() + if include_temporary and message._no_save: + message_data["_no_save"] = True if isinstance(message.content, list): message_data["content"] = [ - part.model_dump() + part.model_dump_for_context() + if include_temporary + else part.model_dump() for part in message.content - if not getattr(part, "_no_save", False) + if include_temporary or not getattr(part, "_no_save", False) ] dumped.append(message_data) if message._checkpoint_after is not None: diff --git a/astrbot/core/agent/response.py b/astrbot/core/agent/response.py index 887f8941a0..4146293865 100644 --- a/astrbot/core/agent/response.py +++ b/astrbot/core/agent/response.py @@ -11,8 +11,33 @@ class AgentResponseData(T.TypedDict): @dataclass class AgentResponse: - type: str - data: AgentResponseData + """One runner output, dispatched by its literal type. + + Display responses carry ``chain``; protocol responses carry event payloads. + ``context.updated`` is an in-memory snapshot for legacy mutation adaptation. + Event IDs identify durable records and must be reused on redelivery. + """ + + type: T.Literal[ + "streaming_delta", + "llm_result", + "err", + "aborted", + "agent_stats", + "tool_call", + "tool_call_result", + "turn.started", + "turn.finished", + "request.started", + "request.finished", + "tool.started", + "tool.finished", + "context.updated", + "message.appended", + "context.rebased", + ] + data: AgentResponseData | dict[str, T.Any] + event_id: str | None = None @dataclass diff --git a/astrbot/core/agent/runners/base.py b/astrbot/core/agent/runners/base.py index d916d0f937..f4c8e88733 100644 --- a/astrbot/core/agent/runners/base.py +++ b/astrbot/core/agent/runners/base.py @@ -20,6 +20,8 @@ class AgentState(Enum): class BaseAgentRunner(T.Generic[TContext]): + conversation_event_capabilities: frozenset[str] = frozenset({"turn", "messages"}) + @abc.abstractmethod async def reset( self, diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index c9787ed6f0..dab13085d9 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -5,7 +5,7 @@ import traceback import typing as T import uuid -from contextlib import suppress +from contextlib import aclosing, suppress from dataclasses import dataclass, field, replace from pathlib import Path @@ -25,7 +25,17 @@ ) from astrbot import logger -from astrbot.core.agent.message import ImageURLPart, TextPart, ThinkPart +from astrbot.core.agent.event_stream import ( + AgentEventStream, + RequestEventRecorder, + request_recorder_kwargs, +) +from astrbot.core.agent.message import ( + ImageURLPart, + TextPart, + ThinkPart, + dump_messages_with_checkpoints, +) from astrbot.core.agent.tool import FunctionTool, ToolSet from astrbot.core.agent.tool_image_cache import tool_image_cache from astrbot.core.exceptions import EmptyModelOutputError @@ -108,6 +118,9 @@ class _ToolExecutionInterrupted(Exception): class ToolLoopAgentRunner(BaseAgentRunner[TContext]): + conversation_event_capabilities = frozenset( + {"turn", "messages", "request", "tool", "context"} + ) TOOL_RESULT_MAX_ESTIMATED_TOKENS = 27_500 TOOL_RESULT_PREVIEW_MAX_ESTIMATED_TOKENS = 7000 EMPTY_OUTPUT_RETRY_ATTEMPTS = 3 @@ -229,9 +242,13 @@ async def reset( request_max_retries: int | None = None, tool_result_overflow_dir: str | None = None, read_tool: FunctionTool | None = None, + turn_id: str | None = None, **kwargs: T.Any, ) -> None: self.req = request + self.turn_id = turn_id or str(uuid.uuid4()) + self._turn_started = False + self._events = AgentEventStream() self.streaming = streaming self.enforce_max_turns = enforce_max_turns self.llm_compress_instruction = llm_compress_instruction @@ -255,6 +272,8 @@ async def reset( llm_compress_provider=self.llm_compress_provider, custom_token_counter=self.custom_token_counter, custom_compressor=self.custom_compressor, + request_event_emitter=self._events.emit, + turn_id=self.turn_id, ) self.request_context_manager = ContextManager( self.request_context_manager_config @@ -497,6 +516,93 @@ async def _await_or_stop( abort_task.cancel() await asyncio.gather(abort_task, return_exceptions=True) + async def _recorded_responses(self, payload: dict, *, streaming: bool): + """Emit provider attempt events without including temporary request bodies. + + Args: + payload: Actual provider call arguments, kept in memory only. + streaming: Whether to consume a streaming provider. + + Yields: + Original provider responses. + """ + recorder = None + final = None + status = "failed" + error_code = "NO_FINAL_RESPONSE" + if self._events.active: + recorder = RequestEventRecorder( + self._events.emit, + { + "turn_id": self.turn_id, + "provider_id": self.provider.provider_config.get("id", ""), + "model": payload.get("model") or self.provider.get_model(), + "attempt_scope": "provider_call", + }, + ) + await recorder.begin() + try: + if streaming: + stream = self.provider.text_chat_stream( + **payload, + **request_recorder_kwargs(self.provider.text_chat_stream, recorder), + ) + try: + while True: + try: + resp = await self._await_or_stop(anext(stream)) + except StopAsyncIteration: + break + if resp is None: + status = "cancelled" + break + if not resp.is_chunk: + final = resp + status = "failed" if resp.role == "err" else "completed" + if recorder: + await recorder.finish( + status, + usage=resp.usage, + error_code="PROVIDER_ERROR" + if status == "failed" + else None, + ) + yield resp + finally: + await self._close_executor(stream) + else: + resp = await self._await_or_stop( + self.provider.text_chat( + **payload, + **request_recorder_kwargs(self.provider.text_chat, recorder), + ) + ) + if resp is None: + status = "cancelled" + else: + final = resp + status = "failed" if resp.role == "err" else "completed" + if recorder: + await recorder.finish( + status, + usage=resp.usage, + error_code="PROVIDER_ERROR" if status == "failed" else None, + ) + yield resp + except (asyncio.CancelledError, GeneratorExit): + status = "cancelled" + raise + except Exception as exc: + error_code = type(exc).__name__ + raise + finally: + if recorder: + await recorder.finish( + status, + usage=final.usage if final else None, + error_code=error_code if status == "failed" else None, + ) + async def _iter_llm_responses( self, *, include_model: bool = True ) -> T.AsyncGenerator[LLMResponse, None]: @@ -512,23 +618,10 @@ async def _iter_llm_responses( if include_model: # For primary provider we keep explicit model selection if provided. payload["model"] = self.req.model - if self.streaming: - stream = self.provider.text_chat_stream(**payload) - try: - while True: - try: - resp = await self._await_or_stop(anext(stream)) # type: ignore - except StopAsyncIteration: - return - if resp is None: - return - yield resp - finally: - await self._close_executor(stream) - else: - resp = await self._await_or_stop(self.provider.text_chat(**payload)) - if resp is not None: - yield resp + async for response in self._recorded_responses( + payload, streaming=self.streaming + ): + yield response async def _iter_llm_responses_with_fallback( self, @@ -794,8 +887,61 @@ def _sanitize_malformed_tool_calls( for tool_name in llm_resp.tools_call_name ] + def _context_response(self, reason="legacy_replace") -> AgentResponse: + """Capture working context for the host's legacy mutation adapter. + + Args: + reason: Why context may have changed, including compaction. + + Returns: + A runtime snapshot; the host stages immutable context events from it. + """ + return AgentResponse( + "context.updated", + { + "messages": dump_messages_with_checkpoints( + [ + m + for i, m in enumerate(self.run_context.messages) + if not (i == 0 and m.role == "system") + ], + include_temporary=True, + ), + "reason": reason, + }, + ) + @override async def step(self): + """Yield execution events and display responses for host consumption. + + Yields: + AgentResponse values. Advance the iterator after processing each value; + close it to cancel pending work when event persistence fails. + """ + if not self._turn_started: + yield AgentResponse( + "turn.started", {"trigger": {"kind": "agent"}}, self.turn_id + ) + self._turn_started = True + yield self._context_response() + async with aclosing(self._events.run(self._step())) as responses: + async for response in responses: + yield response + yield self._context_response() + if self.done(): + yield AgentResponse( + "turn.finished", + { + "turn_id": self.turn_id, + "status": "cancelled" + if self.was_aborted() + else ("failed" if self.state == AgentState.ERROR else "completed"), + }, + str(uuid.uuid4()), + ) + + async def _step(self): """Process a single step of the agent. This method should return the result of the step. """ @@ -825,6 +971,7 @@ async def step(self): yield await self._finalize_aborted_step() return self.run_context.messages = processed_messages + yield self._context_response(reason="compaction") self._simple_print_message_role("[AftCompact]", self.run_context.messages) async for llm_response in self._iter_llm_responses_with_fallback(): @@ -1095,8 +1242,9 @@ async def step_until_done( step_count = 0 while not self.done() and step_count < max_step: step_count += 1 - async for resp in self.step(): - yield resp + async with aclosing(self.step()) as responses: + async for resp in responses: + yield resp # 如果循环结束了但是 agent 还没有完成,说明是达到了 max_step if not self.done(): @@ -1114,8 +1262,9 @@ async def step_until_done( ) ) # 再执行最后一步 - async for resp in self.step(): - yield resp + async with aclosing(self.step()) as responses: + async for resp in responses: + yield resp async def _handle_function_tools( self, @@ -1227,14 +1376,10 @@ def _append_tool_call_result(tool_call_id: str, content: str) -> None: except Exception as e: logger.error(f"Error in on_tool_start hook: {e}", exc_info=True) - executor = self.tool_executor.execute( - tool=func_tool, - run_context=self.run_context, - **valid_params, # 只传递有效的参数 - ) - _final_resp: CallToolResult | None = None - async for resp in self._iter_tool_executor_results(executor): # type: ignore + async for resp in self._recorded_tool_results( + func_tool, func_tool_name, func_tool_id, valid_params + ): if isinstance(resp, CallToolResult): res = resp _final_resp = resp @@ -1441,18 +1586,20 @@ async def _resolve_tool_exec( ) if param_subset.tools and tool_names: contexts = self._build_tool_requery_context(tool_names) - requery_resp = await self._await_or_stop( - self.provider.text_chat( - contexts=self._sanitize_contexts_for_provider(contexts), - func_tool=param_subset, - model=self.req.model, - session_id=self.req.session_id, - extra_user_content_parts=self.req.extra_user_content_parts, - # tool_choice="required", - abort_signal=self._abort_signal, - request_max_retries=self.request_max_retries, - ) - ) + requery_resp = None + async for response in self._recorded_responses( + { + "contexts": self._sanitize_contexts_for_provider(contexts), + "func_tool": param_subset, + "model": self.req.model, + "session_id": self.req.session_id, + "extra_user_content_parts": self.req.extra_user_content_parts, + "abort_signal": self._abort_signal, + "request_max_retries": self.request_max_retries, + }, + streaming=False, + ): + requery_resp = response if requery_resp: llm_resp = requery_resp self._sanitize_malformed_tool_calls(llm_resp) @@ -1471,20 +1618,22 @@ async def _resolve_tool_exec( tool_names, extra_instruction=self.SKILLS_LIKE_REQUERY_REPAIR_INSTRUCTION, ) - repair_resp = await self._await_or_stop( - self.provider.text_chat( - contexts=self._sanitize_contexts_for_provider( + repair_resp = None + async for response in self._recorded_responses( + { + "contexts": self._sanitize_contexts_for_provider( repair_contexts ), - func_tool=param_subset, - model=self.req.model, - session_id=self.req.session_id, - extra_user_content_parts=self.req.extra_user_content_parts, - # tool_choice="required", - abort_signal=self._abort_signal, - request_max_retries=self.request_max_retries, - ) - ) + "func_tool": param_subset, + "model": self.req.model, + "session_id": self.req.session_id, + "extra_user_content_parts": self.req.extra_user_content_parts, + "abort_signal": self._abort_signal, + "request_max_retries": self.request_max_retries, + }, + streaming=False, + ): + repair_resp = response if repair_resp: llm_resp = repair_resp self._sanitize_malformed_tool_calls(llm_resp) @@ -1550,6 +1699,67 @@ async def _close_executor(self, executor: T.Any) -> None: with suppress(asyncio.CancelledError, RuntimeError, StopAsyncIteration): await close_executor() + async def _recorded_tool_results(self, tool, tool_name, tool_call_id, arguments): + """Emit tool execution boundaries around the actual executor. + + Args: + tool: Tool whose executor is created after the started event is handled. + tool_name: Executed tool name. + tool_call_id: Model's call identity. + arguments: Arguments after plugin hooks and validation. + + Yields: + Original tool results. + """ + execution_id = None + result = None + status = "completed" + error_code = None + if self._events.active: + execution_id = str(uuid.uuid4()) + await self._events.emit( + AgentResponse( + "tool.started", + { + "turn_id": self.turn_id, + "tool_call_id": tool_call_id, + "tool_name": tool_name, + "arguments": arguments, + }, + execution_id, + ) + ) + executor = None + try: + executor = self.tool_executor.execute( + tool=tool, run_context=self.run_context, **arguments + ) + async for item in self._iter_tool_executor_results(executor): + if isinstance(item, CallToolResult): + result = item.model_dump(mode="json", exclude_none=True) + if item.isError: + status = "failed" + yield item + except (asyncio.CancelledError, _ToolExecutionInterrupted, GeneratorExit): + status = "cancelled" + raise + except Exception as exc: + status = "failed" + error_code = type(exc).__name__ + raise + finally: + if executor is not None: + await self._close_executor(executor) + if execution_id: + payload = {"execution_id": execution_id, "status": status} + if result is not None: + payload["result"] = result + if error_code: + payload["error"] = {"code": error_code} + await self._events.emit( + AgentResponse("tool.finished", payload, str(uuid.uuid4())) + ) + async def _iter_tool_executor_results( self, executor: T.AsyncGenerator[ToolExecutorResultT, None], @@ -1589,6 +1799,9 @@ async def _next_executor_result() -> ToolExecutorResultT: except StopAsyncIteration: return finally: + if not next_result_task.done(): + next_result_task.cancel() + await asyncio.gather(next_result_task, return_exceptions=True) if not abort_task.done(): abort_task.cancel() with suppress(asyncio.CancelledError): diff --git a/astrbot/core/astr_agent_run_util.py b/astrbot/core/astr_agent_run_util.py index 7e0844ee48..72f211a69b 100644 --- a/astrbot/core/astr_agent_run_util.py +++ b/astrbot/core/astr_agent_run_util.py @@ -3,8 +3,10 @@ import time import traceback from collections.abc import AsyncGenerator +from contextlib import aclosing from astrbot.core import logger +from astrbot.core.agent.conversation_events import ConversationEventWriter from astrbot.core.agent.message import Message from astrbot.core.agent.runners.tool_loop_agent_runner import ToolLoopAgentRunner from astrbot.core.astr_agent_context import AstrAgentContext @@ -120,8 +122,11 @@ async def run_agent( stream_to_general: bool = False, show_reasoning: bool = False, buffer_intermediate_messages: bool = False, + conversation_events: ConversationEventWriter | None = None, ) -> AsyncGenerator[MessageChain | None, None]: step_idx = 0 + if conversation_events is not None: + conversation_events.runtime_context = agent_runner.run_context astr_event = agent_runner.run_context.context.event tool_name_by_call_id: dict[str, str] = {} buffered_llm_chains: list[MessageChain] = [] @@ -153,145 +158,165 @@ async def run_agent( _watch_agent_stop_signal(agent_runner, astr_event), ) try: - async for resp in agent_runner.step(): - if _should_stop_agent(astr_event): - agent_runner.request_stop() - - if resp.type == "aborted": - if can_buffer_llm_result: - merged_chain = _merge_buffered_llm_chains(buffered_llm_chains) - if merged_chain: - astr_event.set_result( - MessageEventResult( - chain=merged_chain.chain, - result_content_type=ResultContentType.LLM_RESULT, - ), + async with aclosing(agent_runner.step()) as responses: + async for resp in responses: + if conversation_events is not None: + if await conversation_events.consume(resp): + continue + if resp.type in { + "context.updated", + "turn.started", + "turn.finished", + "request.started", + "request.finished", + "tool.started", + "tool.finished", + "message.appended", + "context.rebased", + }: + continue + if _should_stop_agent(astr_event): + agent_runner.request_stop() + + if resp.type == "aborted": + if can_buffer_llm_result: + merged_chain = _merge_buffered_llm_chains( + buffered_llm_chains ) - yield merged_chain - astr_event.clear_result() - if not stop_watcher.done(): - stop_watcher.cancel() - try: - await stop_watcher - except asyncio.CancelledError: - pass - astr_event.set_extra("agent_user_aborted", True) - astr_event.set_extra("agent_stop_requested", False) - return - - if _should_stop_agent(astr_event): - continue - - if resp.type == "agent_stats": - if astr_event.get_platform_name() == "webchat": - await astr_event.send(resp.data["chain"]) - continue - - if resp.type == "tool_call_result": - msg_chain = resp.data["chain"] - - astr_event.trace.record( - "agent_tool_result", - tool_result=msg_chain.get_plain_text( - with_other_comps_mark=True - ), - ) + if merged_chain: + astr_event.set_result( + MessageEventResult( + chain=merged_chain.chain, + result_content_type=ResultContentType.LLM_RESULT, + ), + ) + yield merged_chain + astr_event.clear_result() + if not stop_watcher.done(): + stop_watcher.cancel() + try: + await stop_watcher + except asyncio.CancelledError: + pass + astr_event.set_extra("agent_user_aborted", True) + astr_event.set_extra("agent_stop_requested", False) + return + + if _should_stop_agent(astr_event): + continue - if msg_chain.type == "tool_direct_result": - # tool_direct_result 用于标记 llm tool 需要直接发送给用户的内容 - await astr_event.send(msg_chain) + if resp.type == "agent_stats": + if astr_event.get_platform_name() == "webchat": + await astr_event.send(resp.data["chain"]) continue - if astr_event.get_platform_id() == "webchat": - await astr_event.send(msg_chain) - elif show_tool_use and show_tool_call_result: - status_msg = _build_tool_result_status_message( - msg_chain, tool_name_by_call_id - ) - await astr_event.send( - MessageChain(type="tool_call").message(status_msg) + + if resp.type == "tool_call_result": + msg_chain = resp.data["chain"] + + astr_event.trace.record( + "agent_tool_result", + tool_result=msg_chain.get_plain_text( + with_other_comps_mark=True + ), ) - # 对于其他情况,暂时先不处理 - continue - elif resp.type == "tool_call": - if agent_runner.streaming and show_tool_use: - # 向下游平台发送 "break" 分段信号(空 MessageChain,不携带数据)。 - # 平台适配器收到后会关闭当前流式消息,并在后续文本到来时创建新消息。 - # 仅在 show_tool_use 为 True 时才发送:此时紧接着会通过 - # astr_event.send() 独立发送工具状态消息(如"🔨 调用工具: xxx"), - # 需要分段才能保证消息顺序正确。 - # 若 show_tool_use 为 False,不会有独立消息插入,无需分段。 - yield MessageChain(chain=[], type="break") - - tool_info = _extract_chain_json_data(resp.data["chain"]) - astr_event.trace.record( - "agent_tool_call", - tool_name=tool_info if tool_info else "unknown", - ) - _record_tool_call_name(tool_info, tool_name_by_call_id) - if astr_event.get_platform_name() == "webchat": - await astr_event.send(resp.data["chain"]) - elif show_tool_use: - if show_tool_call_result and isinstance(tool_info, dict): - # Delay tool status notification until tool_call_result. + if msg_chain.type == "tool_direct_result": + # tool_direct_result 用于标记 llm tool 需要直接发送给用户的内容 + await astr_event.send(msg_chain) continue - chain = MessageChain(type="tool_call").message( - _build_tool_call_status_message(tool_info) + if astr_event.get_platform_id() == "webchat": + await astr_event.send(msg_chain) + elif show_tool_use and show_tool_call_result: + status_msg = _build_tool_result_status_message( + msg_chain, tool_name_by_call_id + ) + await astr_event.send( + MessageChain(type="tool_call").message(status_msg) + ) + # 对于其他情况,暂时先不处理 + continue + elif resp.type == "tool_call": + if agent_runner.streaming and show_tool_use: + # 向下游平台发送 "break" 分段信号(空 MessageChain,不携带数据)。 + # 平台适配器收到后会关闭当前流式消息,并在后续文本到来时创建新消息。 + # 仅在 show_tool_use 为 True 时才发送:此时紧接着会通过 + # astr_event.send() 独立发送工具状态消息(如"🔨 调用工具: xxx"), + # 需要分段才能保证消息顺序正确。 + # 若 show_tool_use 为 False,不会有独立消息插入,无需分段。 + yield MessageChain(chain=[], type="break") + + tool_info = _extract_chain_json_data(resp.data["chain"]) + astr_event.trace.record( + "agent_tool_call", + tool_name=tool_info if tool_info else "unknown", ) - await astr_event.send(chain) - continue - elif resp.type == "llm_result": - chain = resp.data["chain"] - if chain.type == "reasoning": - # For non-streaming mode, we handle reasoning in astrbot/core/astr_agent_hooks.py. - # For streaming mode, we yield content immediately when received a reasoning chunk but not in here, see below. + _record_tool_call_name(tool_info, tool_name_by_call_id) + + if astr_event.get_platform_name() == "webchat": + await astr_event.send(resp.data["chain"]) + elif show_tool_use: + if show_tool_call_result and isinstance(tool_info, dict): + # Delay tool status notification until tool_call_result. + continue + chain = MessageChain(type="tool_call").message( + _build_tool_call_status_message(tool_info) + ) + await astr_event.send(chain) continue + elif resp.type == "llm_result": + chain = resp.data["chain"] + if chain.type == "reasoning": + # For non-streaming mode, we handle reasoning in astrbot/core/astr_agent_hooks.py. + # For streaming mode, we yield content immediately when received a reasoning chunk but not in here, see below. + continue - if stream_to_general and resp.type == "streaming_delta": - continue + if stream_to_general and resp.type == "streaming_delta": + continue - if ( - resp.type == "err" - and agent_runner.streaming - and not stream_to_general - ): - chain = ( - resp.data.get("chain") if isinstance(resp.data, dict) else None - ) - if not isinstance(chain, MessageChain): - logger.error( - "Agent runner returned an error response without a message chain." - ) - chain = MessageChain().message( - "Error occurred during AI execution." + if ( + resp.type == "err" + and agent_runner.streaming + and not stream_to_general + ): + chain = ( + resp.data.get("chain") + if isinstance(resp.data, dict) + else None ) - yield chain - continue - - if stream_to_general or not agent_runner.streaming: - if can_buffer_llm_result and resp.type == "llm_result": - buffered_llm_chains.append(resp.data["chain"]) + if not isinstance(chain, MessageChain): + logger.error( + "Agent runner returned an error response without a message chain." + ) + chain = MessageChain().message( + "Error occurred during AI execution." + ) + yield chain continue - content_typ = ( - ResultContentType.LLM_RESULT - if resp.type == "llm_result" - else ResultContentType.GENERAL_RESULT - ) - astr_event.set_result( - MessageEventResult( - chain=resp.data["chain"].chain, - result_content_type=content_typ, - ), - ) - yield resp.data["chain"] - astr_event.clear_result() - elif resp.type == "streaming_delta": - chain = resp.data["chain"] - if chain.type == "reasoning" and not show_reasoning: - # display the reasoning content only when configured - continue - yield resp.data["chain"] # MessageChain + if stream_to_general or not agent_runner.streaming: + if can_buffer_llm_result and resp.type == "llm_result": + buffered_llm_chains.append(resp.data["chain"]) + continue + + content_typ = ( + ResultContentType.LLM_RESULT + if resp.type == "llm_result" + else ResultContentType.GENERAL_RESULT + ) + astr_event.set_result( + MessageEventResult( + chain=resp.data["chain"].chain, + result_content_type=content_typ, + ), + ) + yield resp.data["chain"] + astr_event.clear_result() + elif resp.type == "streaming_delta": + chain = resp.data["chain"] + if chain.type == "reasoning" and not show_reasoning: + # display the reasoning content only when configured + continue + yield resp.data["chain"] # MessageChain if can_buffer_llm_result and agent_runner.done(): merged_chain = _merge_buffered_llm_chains(buffered_llm_chains) @@ -351,6 +376,10 @@ async def run_agent( else: astr_event.set_result(MessageEventResult().message(err_msg)) return + finally: + if not stop_watcher.done(): + stop_watcher.cancel() + await asyncio.gather(stop_watcher, return_exceptions=True) async def _watch_agent_stop_signal(agent_runner: AgentRunner, astr_event) -> None: @@ -369,6 +398,7 @@ async def run_live_agent( show_tool_call_result: bool = False, show_reasoning: bool = False, buffer_intermediate_messages: bool = False, + conversation_events: ConversationEventWriter | None = None, ) -> AsyncGenerator[MessageChain | None, None]: """Live Mode 的 Agent 运行器,支持流式 TTS @@ -393,6 +423,7 @@ async def run_live_agent( stream_to_general=False, show_reasoning=show_reasoning, buffer_intermediate_messages=buffer_intermediate_messages, + conversation_events=conversation_events, ): yield chain return @@ -426,6 +457,7 @@ async def run_live_agent( show_tool_call_result, show_reasoning, buffer_intermediate_messages, + conversation_events, ) ) @@ -518,6 +550,7 @@ async def _run_agent_feeder( show_tool_call_result: bool, show_reasoning: bool, buffer_intermediate_messages: bool, + conversation_events: ConversationEventWriter | None = None, ) -> None: """运行 Agent 并将文本输出分句放入队列""" buffer = "" @@ -530,6 +563,7 @@ async def _run_agent_feeder( stream_to_general=False, show_reasoning=show_reasoning, buffer_intermediate_messages=buffer_intermediate_messages, + conversation_events=conversation_events, ): if chain is None: continue diff --git a/astrbot/core/astr_agent_tool_exec.py b/astrbot/core/astr_agent_tool_exec.py index e94f0ce6f9..241237631f 100644 --- a/astrbot/core/astr_agent_tool_exec.py +++ b/astrbot/core/astr_agent_tool_exec.py @@ -6,6 +6,7 @@ import uuid from collections.abc import Sequence from collections.abc import Set as AbstractSet +from contextlib import aclosing import mcp @@ -619,29 +620,46 @@ async def _wake_main_agent_for_background_result( return runner = result.agent_runner - async for _ in runner.step_until_done(agent_max_step): - # agent will send message to user via using tools - pass - llm_resp = runner.get_final_llm_resp() - task_meta = extras.get("background_task_result", {}) - summary_note = ( - f"[BackgroundTask] {summary_name} " - f"(task_id={task_meta.get('task_id', task_id)}) finished. " - f"Result: {task_meta.get('result') or result_text or 'no content'}" - ) - if llm_resp and llm_resp.completion_text: - summary_note += ( - f"I finished the task, here is the result: {llm_resp.completion_text}" + event_writer = result.conversation_events + if event_writer is not None: + event_writer.runtime_context = runner.run_context + status = "failed" + try: + async with aclosing(runner.step_until_done(agent_max_step)) as responses: + async for response in responses: + if event_writer is not None: + await event_writer.consume(response) + llm_resp = runner.get_final_llm_resp() + task_meta = extras.get("background_task_result", {}) + summary_note = ( + f"[BackgroundTask] {summary_name} " + f"(task_id={task_meta.get('task_id', task_id)}) finished. " + f"Result: {task_meta.get('result') or result_text or 'no content'}" ) - await persist_agent_history( - ctx.conversation_manager, - event=cron_event, - req=req, - summary_note=summary_note, - ) - if not llm_resp: - logger.warning("background task agent got no response") - return + if llm_resp and llm_resp.completion_text: + summary_note += f"I finished the task, here is the result: {llm_resp.completion_text}" + await persist_agent_history( + ctx.conversation_manager, + event=cron_event, + req=req, + summary_note=summary_note, + conversation_events=event_writer, + ) + status = ( + "completed" if llm_resp and llm_resp.role == "assistant" else "failed" + ) + if not llm_resp: + logger.warning("background task agent got no response") + return + except asyncio.CancelledError: + status = "cancelled" + raise + finally: + try: + if event_writer is not None: + await event_writer.finish_turn(status) + finally: + cron_event.conversation_events = None @classmethod async def _execute_local( diff --git a/astrbot/core/astr_main_agent.py b/astrbot/core/astr_main_agent.py index 3148d638b6..f82028cc78 100644 --- a/astrbot/core/astr_main_agent.py +++ b/astrbot/core/astr_main_agent.py @@ -12,6 +12,7 @@ from pathlib import Path from astrbot.core import logger +from astrbot.core.agent.conversation_events import ConversationEventWriter from astrbot.core.agent.handoff import HandoffTool from astrbot.core.agent.mcp_client import MCPTool from astrbot.core.agent.message import TextPart @@ -30,7 +31,7 @@ TOOL_CALL_PROMPT_SKILLS_LIKE_MODE, ) from astrbot.core.computer.booters.local import resolve_windows_shell -from astrbot.core.conversation_mgr import Conversation +from astrbot.core.conversation_mgr import Conversation, ConversationManager from astrbot.core.db import BaseDatabase from astrbot.core.message.components import File, Image, Record, Reply, Video from astrbot.core.persona_error_reply import ( @@ -225,6 +226,7 @@ class MainAgentBuildResult: provider_request: ProviderRequest provider: Provider reset_coro: Coroutine | None = None + conversation_events: ConversationEventWriter | None = None def _set_llm_error_message(event: AstrMessageEvent, message: str) -> None: @@ -1716,7 +1718,20 @@ async def build_main_agent( _apply_web_search_citation_prompt(event, req) + event_writer = None + if req.conversation and isinstance( + plugin_context.conversation_manager, ConversationManager + ): + event_writer = await plugin_context.conversation_manager.event_writer( + event.unified_msg_origin, + req.conversation.cid, + expected_revision=req.conversation.revision, + ) + event_writer.request = req + event.conversation_events = event_writer + reset_coro = agent_runner.reset( + turn_id=event.get_extra("turn_id"), provider=provider, request=req, run_context=AgentContextWrapper( @@ -1756,4 +1771,5 @@ async def build_main_agent( provider_request=req, provider=provider, reset_coro=reset_coro if not apply_reset else None, + conversation_events=event_writer, ) diff --git a/astrbot/core/backup/constants.py b/astrbot/core/backup/constants.py index 041fa407bc..c3e5f1337d 100644 --- a/astrbot/core/backup/constants.py +++ b/astrbot/core/backup/constants.py @@ -10,7 +10,8 @@ ChatUIProject, CommandConfig, CommandConflict, - ConversationV2, + ConversationEvent, + ConversationV3, Persona, PersonaFolder, PlatformMessageHistory, @@ -42,7 +43,8 @@ # 主数据库模型类映射 MAIN_DB_MODELS: dict[str, type[SQLModel]] = { "platform_stats": PlatformStat, - "conversations": ConversationV2, + "conversations_v3": ConversationV3, + "conversation_events": ConversationEvent, "personas": Persona, "persona_folders": PersonaFolder, "preferences": Preference, diff --git a/astrbot/core/backup/exporter.py b/astrbot/core/backup/exporter.py index a922375998..ba6bf911b8 100644 --- a/astrbot/core/backup/exporter.py +++ b/astrbot/core/backup/exporter.py @@ -12,7 +12,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any -from sqlalchemy import select +from sqlalchemy import select, text from astrbot.core import logger from astrbot.core.config.default import VERSION @@ -208,6 +208,7 @@ async def _export_main_database(self) -> dict[str, list[dict]]: export_data: dict[str, list[dict]] = {} async with self.main_db.get_db() as session: + await session.execute(text("BEGIN")) for table_name, model_class in MAIN_DB_MODELS.items(): try: result = await session.execute(select(model_class)) diff --git a/astrbot/core/backup/importer.py b/astrbot/core/backup/importer.py index 7d30d27c39..b718b00e4e 100644 --- a/astrbot/core/backup/importer.py +++ b/astrbot/core/backup/importer.py @@ -532,7 +532,7 @@ async def _clear_main_db(self) -> None: """清空主数据库所有表""" async with self.main_db.get_db() as session: async with session.begin(): - for table_name, model_class in MAIN_DB_MODELS.items(): + for table_name, model_class in reversed(MAIN_DB_MODELS.items()): try: await session.execute(delete(model_class)) logger.debug(f"已清空表 {table_name}") @@ -574,10 +574,30 @@ async def _import_main_database( """导入主数据库数据""" imported: dict[str, int] = {} + from sqlalchemy import text + + from astrbot.core.db.po import ConversationV2 + + if data.get("conversations") and data.get("conversations_v3"): + raise ValueError("Backup contains ambiguous V2 and V3 conversations") + async with self.main_db.get_db() as session: async with session.begin(): + await session.execute(text("BEGIN IMMEDIATE")) + await session.execute(text("PRAGMA defer_foreign_keys=ON")) + if "conversations" in data: + connection = await session.connection() + await connection.run_sync( + lambda conn: ConversationV2.__table__.create( + conn, checkfirst=True + ) + ) for table_name, rows in data.items(): - model_class = MAIN_DB_MODELS.get(table_name) + model_class = ( + ConversationV2 + if table_name == "conversations" + else MAIN_DB_MODELS.get(table_name) + ) if not model_class: logger.warning(f"未知的表: {table_name}") continue @@ -586,16 +606,40 @@ async def _import_main_database( count = 0 for row in normalized_rows: try: + row = dict(row) + if ( + table_name == "platform_message_history" + and "llm_checkpoint_id" in row + ): + row["turn_id"] = row.pop("llm_checkpoint_id") + if ( + table_name == "webchat_threads" + and "base_checkpoint_id" in row + ): + row["base_event_id"] = row.pop("base_checkpoint_id") # 转换 datetime 字符串为 datetime 对象 row = self._convert_datetime_fields(row, model_class) obj = model_class(**row) session.add(obj) count += 1 except Exception as e: + if table_name in { + "conversations", + "conversations_v3", + "conversation_events", + "platform_message_history", + "webchat_threads", + }: + raise ValueError( + f"Invalid conversation backup row in {table_name}" + ) from e logger.warning(f"导入记录到 {table_name} 失败: {e}") imported[table_name] = count logger.debug(f"导入表 {table_name}: {count} 条记录") + await session.flush() + if "conversations" in data: + await self.main_db.conversation_store.migrate(session=session) return imported diff --git a/astrbot/core/conversation_mgr.py b/astrbot/core/conversation_mgr.py index c97c3622e5..0b28d9cffc 100644 --- a/astrbot/core/conversation_mgr.py +++ b/astrbot/core/conversation_mgr.py @@ -10,9 +10,21 @@ from deprecated import deprecated from astrbot.core import sp +from astrbot.core.agent.conversation_events import ( + ConversationEventWriter, + PluginConversationEvents, + active_conversation_writer, + active_plugin_id, +) from astrbot.core.agent.message import AssistantMessageSegment, UserMessageSegment from astrbot.core.db import BaseDatabase -from astrbot.core.db.po import Conversation, ConversationV2 +from astrbot.core.db.conversation import ConversationConflictError, ConversationStore +from astrbot.core.db.po import ( + Conversation, + ConversationData, + ConversationRead, + ConversationRevision, +) from astrbot.core.utils.datetime_utils import to_utc_timestamp @@ -27,6 +39,84 @@ def __init__(self, db_helper: BaseDatabase) -> None: # 会话删除回调函数列表(用于级联清理,如知识库配置) self._on_session_deleted_callbacks: list[Callable[[str], Awaitable[None]]] = [] + def get_conversation_events(self) -> PluginConversationEvents: + """Access context writes and private events for the current plugin hook. + + Returns: + Plugin facade bound to the current conversation and plugin identity. + + Raises: + RuntimeError: No conversation-bound plugin hook is running. + """ + writer = active_conversation_writer.get() + plugin_id = active_plugin_id.get() + if writer is None or plugin_id is None: + raise RuntimeError( + "Conversation events require an active conversation hook" + ) + return writer.plugin(plugin_id) + + async def event_writer( + self, + unified_msg_origin: str, + conversation_id: str, + *, + expected_revision: ConversationRevision | None = None, + ) -> ConversationEventWriter: + """Bind custom runners or plugins to an owned conversation revision. + + Args: + unified_msg_origin: Owner UMO resolved by the host. + conversation_id: Public conversation identity. + expected_revision: Revision of the context used to prepare this request. + + Returns: + A branch-bound writer implementing the common event protocol. + + Raises: + ValueError: The conversation is missing or belongs to another UMO. + ConversationConflictError: Context or branch changed during preparation. + """ + active = active_conversation_writer.get() + if ( + active + and active.cid == conversation_id + and active.umo == unified_msg_origin + ): + writer = active + else: + snapshot = await self.db.conversation_store.read(conversation_id) + if snapshot is None or snapshot.conversation.umo != unified_msg_origin: + raise ValueError("Conversation is not accessible") + writer = ConversationEventWriter(self.db.conversation_store, snapshot) + if expected_revision is not None and ( + expected_revision.head_seq != writer.head_seq + or expected_revision.leaf_event_id != writer.leaf_event_id + ): + raise ConversationConflictError( + "Conversation changed while preparing the request" + ) + return writer + + async def fork_conversation( + self, unified_msg_origin: str, source_event_id: str + ) -> str: + """Create a side conversation sharing immutable context ancestry. + + Args: + unified_msg_origin: Owner UMO. + source_event_id: Accessible context node to inherit. + + Returns: + Public identity of the new conversation; selection is unchanged. + """ + conv = await self.db.conversation_store.create( + umo=unified_msg_origin, + platform_id=unified_msg_origin.split(":", 1)[0], + parent_event_id=source_event_id, + ) + return conv.conversation_id + def register_on_session_deleted( self, callback: Callable[[str], Awaitable[None]], @@ -59,34 +149,35 @@ async def _trigger_session_deleted(self, unified_msg_origin: str) -> None: f"会话删除回调执行失败 (session: {unified_msg_origin}): {e}", ) - def _convert_conv_from_v2_to_v1( + def _to_conversation( self, - conv_v2: ConversationV2, + record: ConversationData, include_history: bool = True, ) -> Conversation: - """Convert a ConversationV2 object into the legacy Conversation object. + """Convert a database read result into the plugin-facing conversation. Args: - conv_v2: Database conversation object. + record: Detached database read result or legacy conversation data. include_history: Whether to access and serialize the full history. Returns: Legacy-compatible conversation object. """ - created_ts = to_utc_timestamp(conv_v2.created_at) - updated_ts = to_utc_timestamp(conv_v2.updated_at) + created_ts = to_utc_timestamp(record.created_at) + updated_ts = to_utc_timestamp(record.updated_at) created_at = int(created_ts) if created_ts is not None else 0 updated_at = int(updated_ts) if updated_ts is not None else 0 return Conversation( - platform_id=conv_v2.platform_id, - user_id=conv_v2.user_id, - cid=conv_v2.conversation_id, - history=json.dumps(conv_v2.content or []) if include_history else "[]", - title=conv_v2.title, - persona_id=conv_v2.persona_id, + platform_id=record.platform_id, + user_id=record.user_id, + cid=record.conversation_id, + history=json.dumps(record.content or []) if include_history else "[]", + title=record.title, + persona_id=record.persona_id, created_at=created_at, updated_at=updated_at, - token_usage=conv_v2.token_usage, + token_usage=record.token_usage, + revision=record.revision if isinstance(record, ConversationRead) else None, ) async def new_conversation( @@ -210,7 +301,7 @@ async def get_conversation( conv = await self.db.get_conversation_by_id(cid=conversation_id) conv_res = None if conv: - conv_res = self._convert_conv_from_v2_to_v1(conv) + conv_res = self._to_conversation(conv) return conv_res async def get_conversations( @@ -233,7 +324,7 @@ async def get_conversations( ) convs_res = [] for conv in convs: - conv_res = self._convert_conv_from_v2_to_v1(conv) + conv_res = self._to_conversation(conv) convs_res.append(conv_res) return convs_res @@ -269,7 +360,7 @@ async def get_filtered_conversations( convs, cnt = await self.db.get_filtered_conversations(**query_kwargs) convs_res = [] for conv in convs: - conv_res = self._convert_conv_from_v2_to_v1( + conv_res = self._to_conversation( conv, include_history=include_history, ) @@ -298,6 +389,21 @@ async def update_conversation( # 如果没有提供 conversation_id,则获取当前的 conversation_id = await self.get_curr_conversation_id(unified_msg_origin) if conversation_id: + writer = active_conversation_writer.get() + if writer and writer.cid == conversation_id: + if history is not None: + await writer.save_history(history, token_usage=token_usage) + changes = {} + if title is not None: + changes["title"] = title + if persona_id is not None: + changes["persona_id"] = persona_id + if changes or (history is None and token_usage is not None): + payload = {"changes": changes} + if history is None and token_usage is not None: + payload["token_usage"] = token_usage + await writer.append("conversation.updated", payload) + return await self.db.update_conversation( cid=conversation_id, title=title, @@ -370,24 +476,43 @@ async def add_message_pair( Raises: Exception: If the conversation with the given ID is not found """ + if isinstance(self.db.conversation_store, ConversationStore): + writer = active_conversation_writer.get() + if writer and writer.cid == cid: + messages = [ + m.model_dump() if hasattr(m, "model_dump") else m + for m in (user_message, assistant_message) + ] + await writer.save_history( + [entry["message"] for entry in writer._entries] + messages + ) + return + snapshot = await self.db.conversation_store.read(cid) + if snapshot is None: + raise ValueError(f"Conversation with id {cid} not found") + messages = [ + m.model_dump() if hasattr(m, "model_dump") else m + for m in (user_message, assistant_message) + ] + await self.db.conversation_store.append( + cid, + [ + {"type": "message.appended", "payload": {"message": message}} + for message in messages + ], + expected_head=snapshot.conversation.head_seq, + expected_leaf=snapshot.conversation.leaf_event_id, + ) + return conv = await self.db.get_conversation_by_id(cid=cid) - if not conv: - raise Exception(f"Conversation with id {cid} not found") + if conv is None: + raise ValueError(f"Conversation with id {cid} not found") history = conv.content or [] - if isinstance(user_message, UserMessageSegment): - user_msg_dict = user_message.model_dump() - else: - user_msg_dict = user_message - if isinstance(assistant_message, AssistantMessageSegment): - assistant_msg_dict = assistant_message.model_dump() - else: - assistant_msg_dict = assistant_message - history.append(user_msg_dict) - history.append(assistant_msg_dict) - await self.db.update_conversation( - cid=cid, - content=history, + history.extend( + m.model_dump() if hasattr(m, "model_dump") else m + for m in (user_message, assistant_message) ) + await self.db.update_conversation(cid=cid, content=history) async def get_human_readable_context( self, diff --git a/astrbot/core/cron/manager.py b/astrbot/core/cron/manager.py index 15071257de..412866889e 100644 --- a/astrbot/core/cron/manager.py +++ b/astrbot/core/cron/manager.py @@ -2,6 +2,7 @@ import json import re from collections.abc import Awaitable, Callable +from contextlib import aclosing from datetime import datetime, timezone from typing import TYPE_CHECKING, Any from zoneinfo import ZoneInfo @@ -508,40 +509,57 @@ async def _woke_main_agent( raise RuntimeError("Failed to build main agent for cron job.") runner = result.agent_runner - async for _ in runner.step_until_done(agent_max_step): - # agent will send message to user via using tools - pass - llm_resp = runner.get_final_llm_resp() - if runner.state == AgentState.ERROR: - # The run failed (e.g. malformed function call at max steps) but - # no exception escapes the runner; without this the job was - # recorded as completed with last_error=NULL and the user saw - # only intermediate messages (#9980). - detail = ( - f": {llm_resp.completion_text}" - if llm_resp and llm_resp.completion_text - else "" + event_writer = result.conversation_events + if event_writer is not None: + event_writer.runtime_context = runner.run_context + status = "failed" + try: + async with aclosing(runner.step_until_done(agent_max_step)) as responses: + async for response in responses: + if event_writer is not None: + await event_writer.consume(response) + llm_resp = runner.get_final_llm_resp() + if runner.state == AgentState.ERROR: + # The run failed (e.g. malformed function call at max steps) but + # no exception escapes the runner; without this the job was + # recorded as completed with last_error=NULL and the user saw + # only intermediate messages (#9980). + detail = ( + f": {llm_resp.completion_text}" + if llm_resp and llm_resp.completion_text + else "" + ) + raise RuntimeError(f"Cron agent run ended in ERROR state{detail}") + cron_meta = extras.get("cron_job", {}) if extras else {} + summary_note = ( + f"[CronJob] {cron_meta.get('name') or cron_meta.get('id', 'unknown')}: {cron_meta.get('description', '')} " + f" triggered at {cron_meta.get('run_started_at', 'unknown time')}, " ) - raise RuntimeError(f"Cron agent run ended in ERROR state{detail}") - cron_meta = extras.get("cron_job", {}) if extras else {} - summary_note = ( - f"[CronJob] {cron_meta.get('name') or cron_meta.get('id', 'unknown')}: {cron_meta.get('description', '')} " - f" triggered at {cron_meta.get('run_started_at', 'unknown time')}, " - ) - if llm_resp and llm_resp.role == "assistant": - summary_note += ( - f"I finished this job, here is the result: {llm_resp.completion_text}" + if llm_resp and llm_resp.role == "assistant": + summary_note += f"I finished this job, here is the result: {llm_resp.completion_text}" + + await persist_agent_history( + self.ctx.conversation_manager, + event=cron_event, + req=req, + summary_note=summary_note, + conversation_events=event_writer, ) - - await persist_agent_history( - self.ctx.conversation_manager, - event=cron_event, - req=req, - summary_note=summary_note, - ) - if not llm_resp: - logger.warning("Cron job agent got no response") - return + status = ( + "completed" if llm_resp and llm_resp.role == "assistant" else "failed" + ) + if not llm_resp: + logger.warning("Cron job agent got no response") + return + except asyncio.CancelledError: + status = "cancelled" + raise + finally: + try: + if event_writer is not None: + await event_writer.finish_turn(status) + finally: + cron_event.conversation_events = None __all__ = ["CronJobManager"] diff --git a/astrbot/core/db/__init__.py b/astrbot/core/db/__init__.py index 3b4045441d..b89a9308b0 100644 --- a/astrbot/core/db/__init__.py +++ b/astrbot/core/db/__init__.py @@ -13,7 +13,7 @@ ChatUIProject, CommandConfig, CommandConflict, - ConversationV2, + ConversationRead, CronJob, Persona, PersonaFolder, @@ -50,6 +50,9 @@ def __init__(self) -> None: future=True, connect_args=connect_args, ) + from astrbot.core.db.conversation import ConversationStore + + self.conversation_store = ConversationStore(self) self.AsyncSessionLocal = async_sessionmaker( self.engine, class_=AsyncSession, @@ -129,7 +132,7 @@ async def get_conversations( self, user_id: str | None = None, platform_id: str | None = None, - ) -> list[ConversationV2]: + ) -> list[ConversationRead]: """Get all conversations for a specific user and platform_id(optional). content is not included in the result. @@ -137,7 +140,7 @@ async def get_conversations( ... @abc.abstractmethod - async def get_conversation_by_id(self, cid: str) -> ConversationV2: + async def get_conversation_by_id(self, cid: str) -> ConversationRead | None: """Get a specific conversation by its ID.""" ... @@ -146,7 +149,7 @@ async def get_all_conversations( self, page: int = 1, page_size: int = 20, - ) -> list[ConversationV2]: + ) -> list[ConversationRead]: """Get all conversations with pagination.""" ... @@ -159,7 +162,7 @@ async def get_filtered_conversations( search_query: str = "", include_history: bool = True, **kwargs, - ) -> tuple[list[ConversationV2], int]: + ) -> tuple[list[ConversationRead], int]: """Filter conversations by platform IDs and search text. Args: @@ -192,7 +195,7 @@ async def create_conversation( cid: str | None = None, created_at: datetime.datetime | None = None, updated_at: datetime.datetime | None = None, - ) -> ConversationV2: + ) -> ConversationRead: """Create a new conversation.""" ... @@ -226,10 +229,13 @@ async def insert_platform_message_history( content: dict, sender_id: str | None = None, sender_name: str | None = None, - llm_checkpoint_id: str | None = None, + turn_id: str | None = None, max_messages: int | None = None, + llm_checkpoint_id: str | None = None, ) -> PlatformMessageHistory: """Insert a new platform message history record.""" + if turn_id is None: + turn_id = llm_checkpoint_id ... @abc.abstractmethod @@ -237,9 +243,12 @@ async def update_platform_message_history( self, message_id: int, content: dict | None = None, + turn_id: str | None = None, llm_checkpoint_id: str | None = None, ) -> None: """Update a platform message history record.""" + if turn_id is None: + turn_id = llm_checkpoint_id ... @abc.abstractmethod @@ -299,7 +308,7 @@ async def create_webchat_thread( creator: str, parent_session_id: str, parent_message_id: int, - base_checkpoint_id: str, + base_event_id: str, selected_text: str, ) -> WebChatThread: """Create a WebChat side thread.""" diff --git a/astrbot/core/db/conversation.py b/astrbot/core/db/conversation.py new file mode 100644 index 0000000000..1e75803be4 --- /dev/null +++ b/astrbot/core/db/conversation.py @@ -0,0 +1,1173 @@ +"""Transactional event storage and legacy conversation projections.""" + +import json +import uuid +from contextlib import nullcontext +from copy import deepcopy +from dataclasses import dataclass, field +from datetime import datetime, timezone + +from sqlalchemy import literal +from sqlalchemy.orm.attributes import flag_modified +from sqlmodel import col, delete, select, text, update + +from astrbot.core.db.po import ( + ConversationEvent, + ConversationRead, + ConversationRevision, + ConversationV2, + ConversationV3, + PlatformMessageHistory, + WebChatThread, +) +from astrbot.core.sentinels import NOT_GIVEN + +CONTEXT_TYPES = {"message.appended", "context.rebased"} +EVENT_TYPES = CONTEXT_TYPES | { + "conversation.created", + "conversation.updated", + "turn.started", + "turn.finished", + "request.started", + "request.finished", + "tool.started", + "tool.finished", +} +REBASE_REASONS = {"compaction", "reset", "legacy_replace", "migration", "snapshot"} + + +class ConversationConflictError(RuntimeError): + """The conversation changed after a caller read its revision.""" + + +def same_conversation_owner(source: str, target: str) -> bool: + """Check session ownership for cross-conversation ancestry. + + Args: + source: Source UMO resolved by the host. + target: Target UMO resolved by the host. + + Returns: + Whether the UMOs match or identify WebChat sessions of the same creator. + """ + if source == target: + return True + left, right = source.split(":", 2), target.split(":", 2) + if ( + len(left) != 3 + or len(right) != 3 + or left[0] != "webchat" + or right[0] != "webchat" + ): + return False + a, b = left[2].split("!"), right[2].split("!") + return len(a) == len(b) == 3 and a[0] == b[0] == "webchat" and a[1] == b[1] + + +def persistent_messages(messages: list[dict]) -> list[dict]: + """Project working messages while excluding temporary content. + + Args: + messages: Legacy message dictionaries, including checkpoint markers. + + Returns: + Detached dictionaries eligible for future model context. + """ + result = [] + for original in messages: + if original.get("_no_save") or original.get("role") == "_checkpoint": + continue + message = deepcopy(original) + message.pop("_no_save", None) + if isinstance(message.get("content"), list): + message["content"] = [ + {k: v for k, v in part.items() if k != "_no_save"} + if isinstance(part, dict) + else part + for part in message["content"] + if not isinstance(part, dict) or not part.get("_no_save") + ] + result.append(message) + return result + + +@dataclass +class ConversationSnapshot: + """A detached branch projection and its write revision.""" + + conversation: ConversationV3 + entries: list[dict] = field(default_factory=list) + replay_count: int = 0 + replay_bytes: int = 0 + + @property + def messages(self) -> list[dict]: + """Return a mutable copy without exposing stored payload objects.""" + return deepcopy([entry["message"] for entry in self.entries]) + + +class ConversationStore: + """Store shared conversation events using short SQLite write transactions.""" + + def __init__(self, db): + self.db = db + + async def migrate(self, *, session=None) -> None: + """Atomically migrate legacy rows and WebChat links, then drop the old table. + + Args: + session: Optional import transaction; the caller owns its commit. + """ + owns_transaction = session is None + async with ( + self.db.AsyncSessionLocal() if owns_transaction else nullcontext(session) + ) as session: + if owns_transaction: + await session.execute(text("BEGIN IMMEDIATE")) + exists = ( + await session.execute( + text( + "SELECT 1 FROM sqlite_master WHERE type='table' AND name='conversations'" + ) + ) + ).first() + if not exists: + if owns_transaction: + await session.commit() + return + columns = { + r[1] + for r in ( + await session.execute(text("PRAGMA table_info(conversations)")) + ).all() + } + if "token_usage" not in columns: + await session.execute( + text( + "ALTER TABLE conversations ADD COLUMN token_usage INTEGER NOT NULL DEFAULT 0" + ) + ) + rows = await session.stream_scalars( + select(ConversationV2) + .order_by(ConversationV2.inner_conversation_id) + .execution_options(yield_per=32) + ) + async for old in rows: + if ( + await session.execute( + select(ConversationV3.id).where( + ConversationV3.conversation_id == old.conversation_id + ) + ) + ).first(): + raise ValueError( + "Both V2 and V3 contain this conversation; refusing an ambiguous migration" + ) + conv = ConversationV3( + conversation_id=old.conversation_id, + platform_id=old.platform_id, + umo=old.user_id, + title=old.title, + persona_id=old.persona_id, + created_at=old.created_at, + updated_at=old.updated_at, + head_seq=1, + ) + session.add(conv) + await session.flush() + session.add( + ConversationEvent( + conversation_ref=conv.id, + seq=1, + type="conversation.created", + created_at=old.created_at, + payload={ + "platform_id": old.platform_id, + "umo": old.user_id, + "title": old.title, + "persona_id": old.persona_id, + "token_usage": old.token_usage, + }, + ) + ) + pending = [] + groups = [] + for message in old.content or []: + if message.get("role") == "_checkpoint": + marker = message.get("content", {}) + groups.append( + ( + pending, + marker.get("id") if isinstance(marker, dict) else None, + ) + ) + pending = [] + else: + pending.append(message) + if pending: + groups.append((pending, None)) + entries = [] + for messages, checkpoint in groups: + turn_id = str(uuid.uuid4()) + conv.head_seq += 1 + session.add( + ConversationEvent( + conversation_ref=conv.id, + seq=conv.head_seq, + event_id=turn_id, + type="turn.started", + created_at=old.created_at, + payload={ + "trigger": {"kind": "migration"}, + "base_leaf_event_id": conv.leaf_event_id, + }, + ) + ) + first_user = None + for message in messages: + effective = persistent_messages([message]) + conv.head_seq += 1 + event = ConversationEvent( + conversation_ref=conv.id, + seq=conv.head_seq, + parent_event_id=conv.leaf_event_id, + type="message.appended", + created_at=old.created_at, + payload={ + "turn_id": turn_id, + "message": message, + "include_in_context": bool(effective), + }, + ) + session.add(event) + conv.leaf_event_id = event.event_id + if effective: + entries.append( + {"id": event.event_id, "message": effective[0]} + ) + if first_user is None and message.get("role") == "user": + first_user = event.event_id + if checkpoint: + # Copied side histories may reuse checkpoint IDs. Scope by display session. + display_id = old.user_id.rsplit("!", 1)[-1] + records = ( + ( + await session.execute( + select(PlatformMessageHistory).where( + PlatformMessageHistory.turn_id == checkpoint, + PlatformMessageHistory.user_id == display_id, + ) + ) + ) + .scalars() + .all() + ) + for record in records: + record.turn_id = turn_id + record.context_event_id = ( + first_user + if record.content.get("type") == "user" + else conv.leaf_event_id + ) + session.add(record) + if record.content.get("type") == "bot": + await session.execute( + update(WebChatThread) + .where(WebChatThread.parent_message_id == record.id) + .values(base_event_id=conv.leaf_event_id) + ) + conv.head_seq += 1 + session.add( + ConversationEvent( + conversation_ref=conv.id, + seq=conv.head_seq, + type="turn.finished", + created_at=old.updated_at, + payload={ + "turn_id": turn_id, + "status": "completed", + "migrated": True, + }, + ) + ) + conv.head_seq += 1 + baseline = ConversationEvent( + conversation_ref=conv.id, + seq=conv.head_seq, + parent_event_id=conv.leaf_event_id, + type="context.rebased", + created_at=old.updated_at, + payload={"reason": "migration", "messages": entries}, + ) + session.add(baseline) + conv.leaf_event_id = conv.replay_from_event_id = baseline.event_id + conv.updated_at = old.updated_at + flag_modified(conv, "updated_at") + session.add(conv) + await session.flush() + restored = await self.project(session, conv) + expected = persistent_messages( + [m for m in old.content or [] if m.get("role") != "_checkpoint"] + ) + if restored.messages != expected: + raise ValueError("Conversation migration verification failed") + await rows.close() + await session.flush() + await session.execute(text("DROP TABLE conversations")) + if owns_transaction: + await session.commit() + + async def create( + self, + *, + umo: str, + platform_id: str, + content=None, + title=None, + persona_id=None, + cid=None, + created_at=None, + updated_at=None, + parent_event_id=None, + ) -> ConversationV3: + """Create a conversation, optionally inheriting an authorized branch. + + Args: + umo: Owner session identity. + platform_id: Platform instance. + content: Optional initial legacy messages. + title: Optional title. + persona_id: Optional persona. + cid: Optional caller-supplied public identity. + created_at: Optional migration timestamp. + updated_at: Optional migration timestamp. + parent_event_id: Existing context node owned by the same UMO. + + Returns: + Newly committed metadata. + """ + async with self.db.get_db() as session: + await session.execute(text("BEGIN IMMEDIATE")) + baseline_id = None + if parent_event_id: + parent = ( + await session.execute( + select(ConversationEvent).where( + ConversationEvent.event_id == parent_event_id + ) + ) + ).scalar_one_or_none() + owner = ( + await session.get(ConversationV3, parent.conversation_ref) + if parent + else None + ) + if ( + not owner + or not same_conversation_owner(owner.umo, umo) + or parent.type not in CONTEXT_TYPES + ): + raise ValueError("The parent must be an accessible context event") + inherited = await self.project(session, owner, parent_event_id) + baseline_id = inherited.conversation.replay_from_event_id + conv = ConversationV3( + conversation_id=cid or str(uuid.uuid4()), + umo=umo, + platform_id=platform_id, + title=title, + persona_id=persona_id, + created_at=created_at or datetime.now(timezone.utc), + updated_at=updated_at or created_at or datetime.now(timezone.utc), + leaf_event_id=parent_event_id, + replay_from_event_id=baseline_id, + head_seq=1, + ) + session.add(conv) + await session.flush() + payload = { + "umo": umo, + "platform_id": platform_id, + "title": title, + "persona_id": persona_id, + } + if parent_event_id: + payload["forked_from_event_id"] = parent_event_id + session.add( + ConversationEvent( + conversation_ref=conv.id, + seq=1, + type="conversation.created", + payload=payload, + ) + ) + if content: + for message in content: + if message.get("role") != "_checkpoint" and persistent_messages( + [message] + ) != [message]: + conv.head_seq += 1 + excluded = ConversationEvent( + conversation_ref=conv.id, + seq=conv.head_seq, + parent_event_id=conv.leaf_event_id, + type="message.appended", + payload={ + "message": deepcopy(message), + "include_in_context": False, + }, + ) + session.add(excluded) + conv.leaf_event_id = excluded.event_id + conv.head_seq += 1 + baseline = ConversationEvent( + conversation_ref=conv.id, + seq=conv.head_seq, + parent_event_id=conv.leaf_event_id, + type="context.rebased", + payload={ + "reason": "migration", + "messages": [ + {"id": str(uuid.uuid4()), "message": m} + for m in persistent_messages(content) + ], + }, + ) + session.add(baseline) + conv.leaf_event_id = conv.replay_from_event_id = baseline.event_id + conv.updated_at = updated_at or created_at or conv.updated_at + flag_modified(conv, "updated_at") + await session.commit() + return conv + + async def project( + self, session, conv: ConversationV3, leaf=NOT_GIVEN + ) -> ConversationSnapshot: + """Read only the selected ancestry, stopping at a self-contained rebase. + + Args: + session: Active read or write transaction. + conv: Metadata defining the selected branch. + leaf: Optional alternate branch tip; None selects an empty branch. + + Returns: + Detached effective entries and replay cost. + + Raises: + ValueError: The branch is broken or uses an unsupported schema. + """ + conv = ConversationV3(**conv.model_dump()) + if leaf is not NOT_GIVEN: + conv.leaf_event_id = leaf + snapshot = ConversationSnapshot(conv) + conv.replay_from_event_id = None + if conv.leaf_event_id is None: + return snapshot + events = ConversationEvent.__table__ + path = ( + select( + events.c.event_id, + events.c.parent_event_id, + events.c.type, + events.c.version, + events.c.payload, + literal(0).label("depth"), + ) + .where(events.c.event_id == conv.leaf_event_id) + .cte("context_path", recursive=True) + ) + path = path.union_all( + select( + events.c.event_id, + events.c.parent_event_id, + events.c.type, + events.c.version, + events.c.payload, + (path.c.depth + 1).label("depth"), + ) + .join(path, events.c.event_id == path.c.parent_event_id) + .where(path.c.type != "context.rebased") + ) + stream = await session.stream(select(path).order_by(path.c.depth.desc())) + previous = None + async for row in stream.mappings(): + if row["version"] != 1 or row["type"] not in CONTEXT_TYPES: + raise ValueError("Unsupported context event or version") + if row["type"] == "context.rebased": + snapshot.entries = deepcopy(row["payload"]["messages"]) + conv.replay_from_event_id = row["event_id"] + else: + if row["parent_event_id"] != previous: + raise ValueError("Broken context ancestry") + effective = persistent_messages([row["payload"]["message"]]) + if row["payload"].get("include_in_context", True) and effective: + snapshot.entries.append( + { + "id": row["event_id"], + "message": effective[0], + } + ) + previous = row["event_id"] + snapshot.replay_count += 1 + # Count only the tail: a large baseline must not trigger another snapshot. + if row["type"] != "context.rebased": + snapshot.replay_bytes += len( + json.dumps(row["payload"], ensure_ascii=False).encode() + ) + if previous != conv.leaf_event_id: + raise ValueError("Missing context leaf") + return snapshot + + async def read(self, cid: str) -> ConversationSnapshot | None: + """Read an initialized conversation by its public ID. + + Args: + cid: Public conversation identity. + + Returns: + Detached snapshot, or None when absent. + """ + async with self.db.get_db() as session: + conv = ( + await session.execute( + select(ConversationV3).where(ConversationV3.conversation_id == cid) + ) + ).scalar_one_or_none() + if conv: + return await self.project(session, conv) + return None + + async def read_result( + self, session, conv: ConversationV3, include_history=True + ) -> ConversationRead: + """Read context and its revision into an explicit detached result. + + Args: + session: Current transaction. + conv: V3 metadata. + include_history: Whether to project context. + + Returns: + Legacy-compatible fields and the revision used to project them. + """ + usage = ( + await session.execute( + select(ConversationEvent.payload) + .where( + ConversationEvent.conversation_ref == conv.id, + col(ConversationEvent.type).in_( + ["conversation.created", "conversation.updated"] + ), + ConversationEvent.payload["token_usage"].as_integer().is_not(None), + ) + .order_by(ConversationEvent.seq.desc()) + .limit(1) + ) + ).scalar_one_or_none() + return ConversationRead( + inner_conversation_id=conv.id, + conversation_id=conv.conversation_id, + platform_id=conv.platform_id, + user_id=conv.umo, + title=conv.title, + persona_id=conv.persona_id, + created_at=conv.created_at, + updated_at=conv.updated_at, + token_usage=(usage or {}).get("token_usage", 0), + revision=ConversationRevision(conv.head_seq, conv.leaf_event_id), + content=(await self.project(session, conv)).messages + if include_history + else None, + ) + + async def append( + self, + cid: str, + drafts: list[dict], + *, + expected_head=NOT_GIVEN, + expected_leaf=NOT_GIVEN, + history=None, + reason="legacy_replace", + ) -> list[ConversationEvent]: + """Atomically validate and append a batch, optionally committing legacy history. + + Args: + cid: Public conversation identity. + drafts: Events with type, payload and optional stable event ID/parent. + expected_head: Optional optimistic concurrency revision. + expected_leaf: Optional expected selected branch. + history: Optional complete legacy working history to reconcile. + reason: Reason for a non-append history change. + + Returns: + Committed events, including previously committed identical retries. + + Raises: + ConversationConflictError: A stale revision or conflicting ID was supplied. + ValueError: A payload, reference, or event type is invalid. + """ + drafts = json.loads(json.dumps(drafts, allow_nan=False)) + async with self.db.get_db() as session: + await session.execute(text("BEGIN IMMEDIATE")) + conv = ( + await session.execute( + select(ConversationV3).where(ConversationV3.conversation_id == cid) + ) + ).scalar_one_or_none() + if conv is None: + raise ValueError("Conversation not found") + existing = [] + for draft in drafts: + found = ( + await session.execute( + select(ConversationEvent).where( + ConversationEvent.event_id == draft.get("event_id", "") + ) + ) + ).scalar_one_or_none() + if found: + if ( + found.conversation_ref != conv.id + or found.type != draft["type"] + or found.payload != draft.get("payload", {}) + or found.version != draft.get("version", 1) + or ( + "parent_event_id" in draft + and found.parent_event_id != draft["parent_event_id"] + ) + ): + raise ConversationConflictError( + "Event ID already has different content" + ) + existing.append(found) + if existing: + if len(existing) != len(drafts) or history is not None: + raise ConversationConflictError( + "Cannot partially replay a committed batch" + ) + context_events = [ + event for event in existing if event.type in CONTEXT_TYPES + ] + if context_events: + baseline = await session.get( + ConversationEvent, + (conv.id, max(event.seq for event in existing) + 1), + ) + if ( + baseline + and baseline.type == "context.rebased" + and baseline.payload.get("reason") == "snapshot" + and baseline.parent_event_id == context_events[-1].event_id + ): + existing.append(baseline) + return existing + if (expected_head is not NOT_GIVEN and conv.head_seq != expected_head) or ( + expected_leaf is not NOT_GIVEN and conv.leaf_event_id != expected_leaf + ): + raise ConversationConflictError( + "Conversation changed; reload before committing" + ) + snapshot = None + if history is not None: + for message in history: + if message.get("role") != "_checkpoint" and persistent_messages( + [message] + ) != [message]: + drafts.append( + { + "type": "message.appended", + "payload": { + "message": deepcopy(message), + "include_in_context": False, + }, + } + ) + history = persistent_messages(history) + snapshot = await self.project(session, conv) + before = snapshot.messages + if history != before: + if len(history) >= len(before) and history[: len(before)] == before: + drafts.extend( + {"type": "message.appended", "payload": {"message": m}} + for m in history[len(before) :] + ) + else: + drafts.append( + { + "type": "context.rebased", + "payload": { + "reason": "reset" if not history else reason, + "messages": [ + {"id": str(uuid.uuid4()), "message": m} + for m in history + ], + }, + } + ) + committed = [] + for draft in drafts: + kind = draft["type"] + payload = draft.get("payload", {}) + if kind not in EVENT_TYPES and not kind.startswith("plugin."): + raise ValueError(f"Unknown event type: {kind}") + if draft.get("version", 1) != 1 or not isinstance(payload, dict): + raise ValueError("Unsupported event version or payload") + if kind == "conversation.created": + raise ValueError("Use create() to create a conversation") + parent_id = ( + draft.get("parent_event_id", conv.leaf_event_id) + if kind in CONTEXT_TYPES + else None + ) + if ( + kind not in CONTEXT_TYPES + and draft.get("parent_event_id") is not None + ): + raise ValueError("Only context events have context parents") + if parent_id: + parent = ( + await session.execute( + select(ConversationEvent).where( + ConversationEvent.event_id == parent_id + ) + ) + ).scalar_one_or_none() + owner = ( + await session.get(ConversationV3, parent.conversation_ref) + if parent + else None + ) + if ( + not owner + or not same_conversation_owner(owner.umo, conv.umo) + or parent.type not in CONTEXT_TYPES + ): + raise ValueError("Parent must be an accessible context event") + if kind == "message.appended": + message = payload.get("message") + if not isinstance(message, dict) or not isinstance( + message.get("role"), str + ): + raise ValueError("A message requires a role") + if message["role"] == "_checkpoint": + raise ValueError("Checkpoint markers are not model messages") + if not isinstance(payload.get("include_in_context", True), bool): + raise ValueError("include_in_context must be boolean") + elif kind == "context.rebased": + if payload.get("reason") not in REBASE_REASONS or not isinstance( + payload.get("messages"), list + ): + raise ValueError( + "A rebase requires a reason and complete messages" + ) + ids = set() + for entry in payload["messages"]: + if ( + not isinstance(entry, dict) + or not isinstance(entry.get("id"), str) + or entry["id"] in ids + ): + raise ValueError("Rebase message identities must be unique") + ids.add(entry["id"]) + message = entry.get("message") + if ( + not isinstance(message, dict) + or "role" not in message + or persistent_messages([message]) != [message] + ): + raise ValueError("Invalid or temporary rebase message") + elif kind == "conversation.updated": + changes = payload.get("changes", {}) + if set(changes) - {"title", "persona_id"}: + raise ValueError("Unsupported conversation metadata change") + for key, value in changes.items(): + if value is not None and not isinstance(value, str): + raise ValueError("Metadata must be text or null") + setattr(conv, key, value) + elif kind in {"turn.finished", "request.finished", "tool.finished"}: + if payload.get("status") not in { + "completed", + "failed", + "cancelled", + }: + raise ValueError("Invalid execution status") + key = { + "turn.finished": "turn_id", + "request.finished": "request_id", + "tool.finished": "execution_id", + }[kind] + start = ( + await session.execute( + select(ConversationEvent).where( + ConversationEvent.event_id == payload.get(key, "") + ) + ) + ).scalar_one_or_none() + if ( + not start + or start.conversation_ref != conv.id + or start.type != kind.replace("finished", "started") + ): + raise ValueError("Missing matching execution start") + if kind == "request.finished" and "usage" in payload: + usage = payload["usage"] + if not isinstance(usage, dict) or any( + type(value) is not int or value < 0 + for value in usage.values() + ): + raise ValueError( + "Token usage must contain nonnegative integers" + ) + if usage.get("cached_input_tokens", 0) > usage.get( + "input_tokens", 0 + ): + raise ValueError( + "Cached input tokens are a subset of input tokens" + ) + duplicate = ( + await session.execute( + select(ConversationEvent.event_id) + .where( + ConversationEvent.conversation_ref == conv.id, + ConversationEvent.type == kind, + ConversationEvent.seq > start.seq, + ConversationEvent.payload[key].as_string() + == payload[key], + ) + .limit(1) + ) + ).first() + if duplicate: + raise ConversationConflictError("Execution already finished") + elif kind in {"request.started", "tool.started"}: + turn = ( + await session.execute( + select(ConversationEvent).where( + ConversationEvent.event_id == payload.get("turn_id", "") + ) + ) + ).scalar_one_or_none() + if ( + not turn + or turn.conversation_ref != conv.id + or turn.type != "turn.started" + ): + raise ValueError("Missing parent turn") + if kind == "tool.started" and ( + not isinstance(payload.get("tool_name"), str) + or not isinstance(payload.get("tool_call_id"), str) + or not isinstance(payload.get("arguments"), dict) + ): + raise ValueError( + "A tool execution requires a name, call ID and arguments" + ) + conv.head_seq += 1 + event = ConversationEvent( + conversation_ref=conv.id, + seq=conv.head_seq, + event_id=draft.get("event_id") or str(uuid.uuid4()), + parent_event_id=parent_id, + type=kind, + payload=payload, + ) + session.add(event) + await session.flush() + if kind in CONTEXT_TYPES and payload.get("turn_id"): + role = ( + payload["message"].get("role") + if kind == "message.appended" + else payload["messages"][-1]["message"].get("role") + if payload["messages"] + else None + ) + if role in {"user", "assistant"}: + query = update(PlatformMessageHistory).where( + PlatformMessageHistory.turn_id == payload["turn_id"], + PlatformMessageHistory.content["type"].as_string() + == ("user" if role == "user" else "bot"), + ) + if role == "user": + query = query.where( + PlatformMessageHistory.context_event_id.is_(None) + ) + await session.execute( + query.values(context_event_id=event.event_id) + ) + committed.append(event) + if kind in CONTEXT_TYPES: + conv.leaf_event_id = event.event_id + conv.replay_from_event_id = ( + event.event_id if kind == "context.rebased" else None + ) + if committed: + # Snapshots bound replay cost without changing model-visible messages. + if any(e.type in CONTEXT_TYPES for e in committed): + snapshot = await self.project(session, conv) + conv.replay_from_event_id = ( + snapshot.conversation.replay_from_event_id + ) + if snapshot.replay_count > 256 or ( + snapshot.replay_count > 1 + and snapshot.replay_bytes > 4 * 1024 * 1024 + ): + conv.head_seq += 1 + baseline = ConversationEvent( + conversation_ref=conv.id, + seq=conv.head_seq, + type="context.rebased", + parent_event_id=conv.leaf_event_id, + payload={ + "reason": "snapshot", + "messages": snapshot.entries, + }, + ) + session.add(baseline) + conv.leaf_event_id = conv.replay_from_event_id = ( + baseline.event_id + ) + committed.append(baseline) + conv.updated_at = datetime.now(timezone.utc) + session.add(conv) + await session.commit() + return committed + + async def events( + self, + cid: str, + *, + after_seq=0, + before_seq=None, + event_type=None, + limit=100, + max_bytes=4 * 1024 * 1024, + descending=False, + ) -> list[ConversationEvent]: + """Read an indexed event page without loading the conversation projection. + + Args: + cid: Public conversation identity. + after_seq: Exclusive lower cursor. + before_seq: Optional exclusive upper cursor. + event_type: Optional exact event type. + limit: Maximum events, from 1 to 1000. + max_bytes: Soft byte budget; one oversized event is returned alone. + descending: Whether to read newest first. + + Returns: + Detached events; continue using the last returned seq. + """ + if not 1 <= limit <= 1000 or max_bytes <= 0: + raise ValueError("Invalid event page budget") + async with self.db.get_db() as session: + query = ( + select(ConversationEvent) + .join( + ConversationV3, + ConversationV3.id == ConversationEvent.conversation_ref, + ) + .where( + ConversationV3.conversation_id == cid, + ConversationEvent.seq > after_seq, + ) + ) + if before_seq is not None: + query = query.where(ConversationEvent.seq < before_seq) + if event_type is not None: + query = query.where(ConversationEvent.type == event_type) + query = query.order_by( + ConversationEvent.seq.desc() + if descending + else ConversationEvent.seq.asc() + ).limit(limit) + result, size = [], 0 + stream = await session.stream_scalars(query.execution_options(yield_per=1)) + try: + async for event in stream: + weight = len(json.dumps(event.payload, ensure_ascii=False).encode()) + if result and size + weight > max_bytes: + break + result.append(ConversationEvent(**deepcopy(event.model_dump()))) + size += weight + finally: + await stream.close() + return result + + async def select_branch( + self, cid: str, leaf: str | None, *, expected_head, expected_leaf + ) -> None: + """Select an existing branch atomically without appending an event. + + Args: + cid: Target conversation. + leaf: Accessible context node or None. + expected_head: Revision read by the caller. + expected_leaf: Previously selected tip. + """ + async with self.db.get_db() as session: + await session.execute(text("BEGIN IMMEDIATE")) + conv = ( + await session.execute( + select(ConversationV3).where(ConversationV3.conversation_id == cid) + ) + ).scalar_one() + if conv.head_seq != expected_head or conv.leaf_event_id != expected_leaf: + raise ConversationConflictError( + "Conversation changed before branch selection" + ) + if leaf: + target = ( + await session.execute( + select(ConversationEvent).where( + ConversationEvent.event_id == leaf + ) + ) + ).scalar_one_or_none() + owner = ( + await session.get(ConversationV3, target.conversation_ref) + if target + else None + ) + if ( + not owner + or not same_conversation_owner(owner.umo, conv.umo) + or target.type not in CONTEXT_TYPES + ): + raise ValueError("Branch is not accessible") + snapshot = await self.project(session, conv, leaf) + conv.leaf_event_id = leaf + conv.replay_from_event_id = snapshot.conversation.replay_from_event_id + conv.updated_at = datetime.now(timezone.utc) + session.add(conv) + await session.commit() + + async def rewind_webchat( + self, cid: str, message_id: int, *, expected_head, expected_leaf, content=None + ) -> PlatformMessageHistory: + """Fork an edited/retried turn and atomically update its display projection. + + Args: + cid: Conversation already authorized by the WebChat service. + message_id: Active user or assistant display record. + expected_head: Metadata revision read by the service. + expected_leaf: Selected tip read by the service. + content: Optional edited user display content. + + Returns: + New user display record for the replacement turn. + """ + async with self.db.get_db() as session: + await session.execute(text("BEGIN IMMEDIATE")) + conv = ( + await session.execute( + select(ConversationV3).where(ConversationV3.conversation_id == cid) + ) + ).scalar_one() + if conv.head_seq != expected_head or conv.leaf_event_id != expected_leaf: + raise ConversationConflictError("Conversation changed before editing") + target = await session.get(PlatformMessageHistory, message_id) + if not target or not target.is_active or not target.turn_id: + raise ValueError("Message is not an active linked turn") + source = ( + await session.execute( + select(PlatformMessageHistory) + .where( + PlatformMessageHistory.turn_id == target.turn_id, + PlatformMessageHistory.platform_id == target.platform_id, + PlatformMessageHistory.user_id == target.user_id, + PlatformMessageHistory.content["type"].as_string() == "user", + ) + .order_by(PlatformMessageHistory.id) + .limit(1) + ) + ).scalar_one_or_none() + turn = ( + await session.execute( + select(ConversationEvent).where( + ConversationEvent.event_id == target.turn_id, + ConversationEvent.type == "turn.started", + ConversationEvent.conversation_ref == conv.id, + ) + ) + ).scalar_one_or_none() + if source is None or turn is None: + raise ValueError("The original user turn is unavailable") + parent = turn.payload.get("base_leaf_event_id") + snapshot = await self.project(session, conv, parent) + conv.leaf_event_id = parent + conv.replay_from_event_id = snapshot.conversation.replay_from_event_id + conv.updated_at = datetime.now(timezone.utc) + # Old display records remain available to existing side threads. + await session.execute( + update(PlatformMessageHistory) + .where( + PlatformMessageHistory.platform_id == source.platform_id, + PlatformMessageHistory.user_id == source.user_id, + PlatformMessageHistory.id >= source.id, + ) + .values(is_active=False) + ) + replacement = PlatformMessageHistory( + platform_id=source.platform_id, + user_id=source.user_id, + sender_id=source.sender_id, + sender_name=source.sender_name, + content=deepcopy(content if content is not None else source.content), + turn_id=str(uuid.uuid4()), + ) + session.add_all([conv, replacement]) + await session.commit() + return replacement + + async def delete(self, *, cid=None, umo=None) -> None: + """Delete selected conversations only when no surviving branch references them. + + Args: + cid: Optional single public identity. + umo: Optional owner whose conversations are deleted together. + + Raises: + ValueError: A surviving conversation depends on the selected history. + """ + async with self.db.get_db() as session: + await session.execute(text("BEGIN IMMEDIATE")) + query = ( + select(ConversationV3.id).where(ConversationV3.conversation_id == cid) + if cid + else select(ConversationV3.id).where(ConversationV3.umo == umo) + ) + ids = list((await session.execute(query)).scalars()) + if ids: + owned = select(ConversationEvent.event_id).where( + col(ConversationEvent.conversation_ref).in_(ids) + ) + child = ( + await session.execute( + select(ConversationEvent.event_id) + .where( + ~col(ConversationEvent.conversation_ref).in_(ids), + col(ConversationEvent.parent_event_id).in_(owned), + ) + .limit(1) + ) + ).first() + leaf = ( + await session.execute( + select(ConversationV3.id) + .where( + ~col(ConversationV3.id).in_(ids), + col(ConversationV3.leaf_event_id).in_(owned), + ) + .limit(1) + ) + ).first() + if child or leaf: + raise ValueError("Conversation is referenced by another branch") + await session.execute( + delete(ConversationEvent).where( + col(ConversationEvent.conversation_ref).in_(ids) + ) + ) + await session.execute( + delete(ConversationV3).where(col(ConversationV3.id).in_(ids)) + ) + await session.commit() diff --git a/astrbot/core/db/migration/migra_3_to_4.py b/astrbot/core/db/migration/migra_3_to_4.py index 727d97b29b..1fce6ca166 100644 --- a/astrbot/core/db/migration/migra_3_to_4.py +++ b/astrbot/core/db/migration/migra_3_to_4.py @@ -7,7 +7,7 @@ from astrbot.api import logger, sp from astrbot.core.config import AstrBotConfig from astrbot.core.config.default import DB_PATH -from astrbot.core.db.po import ConversationV2, PlatformMessageHistory +from astrbot.core.db.po import PlatformMessageHistory from astrbot.core.platform.astr_message_event import MessageSesion from .. import BaseDatabase @@ -53,48 +53,46 @@ async def migration_conversation_table( ) logger.info(f"迁移 {total_cnt} 条旧的会话数据到新的表中...") - async with db_helper.get_db() as dbsession: - dbsession: AsyncSession - async with dbsession.begin(): - for idx, conversation in enumerate(conversations): - if total_cnt > 0 and (idx + 1) % max(1, total_cnt // 10) == 0: - progress = int((idx + 1) / total_cnt * 100) - if progress % 10 == 0: - logger.info(f"进度: {progress}% ({idx + 1}/{total_cnt})") - try: - conv = db_helper_v3.get_conversation_by_user_id( - user_id=conversation.get("user_id", "unknown"), - cid=conversation.get("cid", "unknown"), - ) - if not conv: - logger.info( - f"未找到该条旧会话对应的具体数据: {conversation}, 跳过。", - ) - continue - if ":" not in conv.user_id: - continue - session = MessageSesion.from_str(session_str=conv.user_id) - platform_id = get_platform_id( - platform_id_map, - session.platform_name, - ) - session.platform_id = platform_id # 更新平台名称为新的 ID - conv_v2 = ConversationV2( - user_id=str(session), - content=json.loads(conv.history) if conv.history else [], - platform_id=platform_id, - title=conv.title, - persona_id=conv.persona_id, - conversation_id=conv.cid, - created_at=datetime.datetime.fromtimestamp(conv.created_at), - updated_at=datetime.datetime.fromtimestamp(conv.updated_at), - ) - dbsession.add(conv_v2) - except Exception as e: - logger.error( - f"迁移旧会话 {conversation.get('cid', 'unknown')} 失败: {e}", - exc_info=True, - ) + for idx, conversation in enumerate(conversations): + if total_cnt > 0 and (idx + 1) % max(1, total_cnt // 10) == 0: + progress = int((idx + 1) / total_cnt * 100) + if progress % 10 == 0: + logger.info(f"进度: {progress}% ({idx + 1}/{total_cnt})") + try: + conv = db_helper_v3.get_conversation_by_user_id( + user_id=conversation.get("user_id", "unknown"), + cid=conversation.get("cid", "unknown"), + ) + if not conv: + logger.info( + f"未找到该条旧会话对应的具体数据: {conversation}, 跳过。", + ) + continue + if ":" not in conv.user_id: + continue + session = MessageSesion.from_str(session_str=conv.user_id) + platform_id = get_platform_id( + platform_id_map, + session.platform_name, + ) + session.platform_id = platform_id # 更新平台名称为新的 ID + if await db_helper.get_conversation_by_id(conv.cid): + continue + await db_helper.create_conversation( + user_id=str(session), + content=json.loads(conv.history) if conv.history else [], + platform_id=platform_id, + title=conv.title, + persona_id=conv.persona_id, + cid=conv.cid, + created_at=datetime.datetime.fromtimestamp(conv.created_at), + updated_at=datetime.datetime.fromtimestamp(conv.updated_at), + ) + except Exception as e: + logger.error( + f"迁移旧会话 {conversation.get('cid', 'unknown')} 失败: {e}", + exc_info=True, + ) logger.info(f"成功迁移 {total_cnt} 条旧的会话数据到新表。") diff --git a/astrbot/core/db/migration/migra_token_usage.py b/astrbot/core/db/migration/migra_token_usage.py index 76bf8ce01c..cfdf66b8f3 100644 --- a/astrbot/core/db/migration/migra_token_usage.py +++ b/astrbot/core/db/migration/migra_token_usage.py @@ -35,7 +35,7 @@ async def migrate_token_usage(db_helper: BaseDatabase) -> None: columns = result.fetchall() column_names = [col[1] for col in columns] - if "token_usage" in column_names: + if not column_names or "token_usage" in column_names: logger.info("token_usage 列已存在,跳过迁移") await sp.put_async( "global", "global", "migration_done_token_usage_1", True diff --git a/astrbot/core/db/migration/migra_webchat_session.py b/astrbot/core/db/migration/migra_webchat_session.py index 5775cbe017..ae2ad3acf2 100644 --- a/astrbot/core/db/migration/migra_webchat_session.py +++ b/astrbot/core/db/migration/migra_webchat_session.py @@ -13,7 +13,7 @@ from sqlmodel import col from astrbot.core.db import BaseDatabase -from astrbot.core.db.po import ConversationV2, PlatformMessageHistory, PlatformSession +from astrbot.core.db.po import ConversationV3, PlatformMessageHistory, PlatformSession async def migrate_webchat_session(db_helper: BaseDatabase) -> None: @@ -55,8 +55,8 @@ async def migrate_webchat_session(db_helper: BaseDatabase) -> None: for user_id, _, _, _ in webchat_users ] conv_query = select( - col(ConversationV2.user_id), col(ConversationV2.title) - ).where(col(ConversationV2.user_id).in_(user_ids_to_query)) + col(ConversationV3.umo), col(ConversationV3.title) + ).where(col(ConversationV3.umo).in_(user_ids_to_query)) conv_result = await session.execute(conv_query) # 创建 user_id -> title 的映射字典 title_map = { diff --git a/astrbot/core/db/po.py b/astrbot/core/db/po.py index 366a07292d..ed488925e8 100644 --- a/astrbot/core/db/po.py +++ b/astrbot/core/db/po.py @@ -4,7 +4,7 @@ from typing import TypedDict from deprecated import deprecated -from sqlalchemy import Index, desc +from sqlalchemy import BigInteger, CheckConstraint, Index, desc from sqlmodel import JSON, Field, SQLModel, Text, UniqueConstraint @@ -64,8 +64,16 @@ class ProviderStat(TimestampMixin, SQLModel, table=True): time_to_first_token: float = Field(default=0.0, nullable=False) -class ConversationV2(TimestampMixin, SQLModel, table=True): - __tablename__: str = "conversations" +@dataclass(frozen=True, slots=True) +class ConversationRevision: + """Committed event position and selected branch observed by a reader.""" + + head_seq: int + leaf_event_id: str | None + + +class ConversationData(TimestampMixin): + """Conversation fields retained by the legacy plugin-facing read API.""" inner_conversation_id: int | None = Field( default=None, @@ -90,6 +98,18 @@ class ConversationV2(TimestampMixin, SQLModel, table=True): when 0, will use estimated token counter. """ + +class ConversationRead(ConversationData): + """Detached read result; content is None when history was not requested.""" + + revision: ConversationRevision + + +class ConversationV2(ConversationData, table=True): + """Legacy table mapping used only to migrate existing data and backups.""" + + __tablename__: str = "conversations" + __table_args__ = ( Index( "ix_conversations_created_at_inner_id", @@ -109,6 +129,53 @@ class ConversationV2(TimestampMixin, SQLModel, table=True): ) +class ConversationV3(TimestampMixin, SQLModel, table=True): + """Current metadata and branch pointers for an event-backed conversation.""" + + __tablename__: str = "conversations_v3" + + id: int | None = Field(default=None, primary_key=True) + conversation_id: str = Field(default_factory=lambda: str(uuid.uuid4()), unique=True) + platform_id: str + umo: str = Field(index=True) + title: str | None = None + persona_id: str | None = None + head_seq: int = Field(default=0, sa_type=BigInteger) + leaf_event_id: str | None = None + replay_from_event_id: str | None = None + + __table_args__ = ( + Index("ix_conversations_v3_created_id", "created_at", "id"), + Index( + "ix_conversations_v3_platform_created_id", "platform_id", "created_at", "id" + ), + ) + + +class ConversationEvent(SQLModel, table=True): + """An immutable context or execution record; seq is conversation-local.""" + + __tablename__: str = "conversation_events" + + conversation_ref: int = Field(foreign_key="conversations_v3.id", primary_key=True) + seq: int = Field(primary_key=True, sa_type=BigInteger) + event_id: str = Field(default_factory=lambda: str(uuid.uuid4()), unique=True) + parent_event_id: str | None = Field( + default=None, foreign_key="conversation_events.event_id" + ) + type: str + version: int = Field(default=1) + payload: dict = Field(default_factory=dict, sa_type=JSON) + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) + + __table_args__ = ( + CheckConstraint("seq > 0"), + CheckConstraint("version > 0"), + Index("ix_conversation_events_type_seq", "conversation_ref", "type", "seq"), + Index("ix_conversation_events_parent", "parent_event_id"), + ) + + class PersonaFolder(TimestampMixin, SQLModel, table=True): """Persona 文件夹,支持递归层级结构。 @@ -257,7 +324,14 @@ class PlatformMessageHistory(TimestampMixin, SQLModel, table=True): default=None, ) # Name of the sender in the platform content: dict = Field(sa_type=JSON, nullable=False) # a message chain list - llm_checkpoint_id: str | None = Field(default=None, index=True) + turn_id: str | None = Field(default=None, index=True) + context_event_id: str | None = Field(default=None, index=True) + is_active: bool = Field(default=True) + + @property + def llm_checkpoint_id(self) -> str | None: + """Expose the old plugin-facing association name during migration.""" + return self.turn_id __table_args__ = ( Index( @@ -288,7 +362,7 @@ class WebChatThread(TimestampMixin, SQLModel, table=True): creator: str = Field(nullable=False, index=True) parent_session_id: str = Field(nullable=False, index=True) parent_message_id: int = Field(nullable=False, index=True) - base_checkpoint_id: str = Field(nullable=False, index=True) + base_event_id: str = Field(nullable=False, index=True) selected_text: str = Field(sa_type=Text, nullable=False) __table_args__ = ( @@ -576,6 +650,8 @@ class Conversation: updated_at: int = 0 token_usage: int = 0 """对话的总 token 数量。AstrBot 会保留最近一次 LLM 请求返回的总 token 数,方便统计。token_usage 可能为 0,表示未知。""" + revision: ConversationRevision | None = None + """Read revision, absent for conversations constructed by legacy plugins.""" class Personality(TypedDict): diff --git a/astrbot/core/db/sqlite.py b/astrbot/core/db/sqlite.py index b644ef3312..c89e5d6388 100644 --- a/astrbot/core/db/sqlite.py +++ b/astrbot/core/db/sqlite.py @@ -8,11 +8,10 @@ from pathlib import Path from deprecated import deprecated -from sqlalchemy import CursorResult, Row, case, literal, not_ +from sqlalchemy import CursorResult, Row, bindparam, case, event, literal, not_ from sqlalchemy.dialects.sqlite import dialect as sqlite_dialect from sqlalchemy.dialects.sqlite import insert as sqlite_insert from sqlalchemy.ext.asyncio import AsyncSession -from sqlalchemy.orm import defer from sqlmodel import col, delete, desc, func, or_, select, text, update from astrbot.core.db import BaseDatabase @@ -22,7 +21,8 @@ ChatUIProject, CommandConfig, CommandConflict, - ConversationV2, + ConversationEvent, + ConversationV3, CronJob, Persona, PersonaFolder, @@ -65,7 +65,7 @@ def _webchat_session_title_match(keyword: str): select(1) .where(col(PlatformSession.display_name).ilike(f"%{keyword}%")) .where( - col(ConversationV2.user_id).like( + col(ConversationV3.umo).like( literal("%!").concat(col(PlatformSession.session_id)), ) ) @@ -73,17 +73,72 @@ def _webchat_session_title_match(keyword: str): ) +def _conversation_content_match(keyword: str): + """Search only effective branch payloads, including inherited rebase content. + + Args: + keyword: User-supplied literal search substring. + + Returns: + A correlated, parameterized SQLite expression. + """ + return text("""EXISTS ( + WITH RECURSIVE branch(event_id, parent_event_id, type, payload) AS ( + SELECT event_id, parent_event_id, type, payload FROM conversation_events + WHERE event_id = conversations_v3.leaf_event_id + UNION ALL + SELECT e.event_id, e.parent_event_id, e.type, e.payload + FROM conversation_events e JOIN branch b ON e.event_id = b.parent_event_id + WHERE b.type != 'context.rebased' + ) + SELECT 1 FROM branch WHERE + (type = 'context.rebased' OR + (type = 'message.appended' AND coalesce(json_extract(payload, '$.include_in_context'), 1))) + AND (json_extract(payload, '$.message') LIKE :pattern + OR json_extract(payload, '$.messages') LIKE :pattern + OR json_extract(payload, '$.message') LIKE :escaped + OR json_extract(payload, '$.messages') LIKE :escaped) + )""").bindparams( + bindparam("pattern", f"%{keyword}%", unique=True), + bindparam( + "escaped", f"%{json.dumps(keyword, ensure_ascii=True)[1:-1]}%", unique=True + ), + ) + + class SQLiteDatabase(BaseDatabase): def __init__(self, db_path: str) -> None: self.db_path = db_path self.DATABASE_URL = f"sqlite+aiosqlite:///{db_path}" self.inited = False + self._initialize_lock = asyncio.Lock() super().__init__() + @event.listens_for(self.engine.sync_engine, "connect") + def enable_foreign_keys(connection, _record): + """Enforce event ownership and ancestry on every pooled connection.""" + cursor = connection.cursor() + cursor.execute("PRAGMA foreign_keys=ON") + cursor.close() + async def initialize(self) -> None: + """Serialize schema migration before allowing application traffic.""" + async with self._initialize_lock: + await self._initialize_schema() + + async def _initialize_schema(self) -> None: """Initialize the database by creating tables if they do not exist.""" async with self.engine.begin() as conn: - await conn.run_sync(SQLModel.metadata.create_all) + await conn.run_sync( + lambda sync_conn: SQLModel.metadata.create_all( + sync_conn, + tables=[ + table + for name, table in SQLModel.metadata.tables.items() + if name != "conversations" + ], + ) + ) await conn.execute(text("PRAGMA journal_mode=WAL")) await conn.execute(text("PRAGMA busy_timeout=30000")) await conn.execute(text("PRAGMA synchronous=NORMAL")) @@ -97,32 +152,12 @@ async def initialize(self) -> None: await self._ensure_persona_custom_error_message_column(conn) await self._ensure_platform_message_history_checkpoint_column(conn) await self._ensure_chatui_project_workspace_columns(conn) - await self._ensure_conversation_indexes(conn) # The table-level unique constraint already provides an index for UMO # lookups. Older schemas also created this redundant explicit index. await conn.execute(text("DROP INDEX IF EXISTS ix_umo_aliases_umo")) await conn.commit() - - async def _ensure_conversation_indexes(self, conn) -> None: - """Create indexes used by the dashboard conversation list. - - Args: - conn: Active SQLAlchemy connection used during SQLite initialization. - """ - await conn.execute( - text( - "CREATE INDEX IF NOT EXISTS " - "ix_conversations_created_at_inner_id " - "ON conversations (created_at DESC, inner_conversation_id DESC)" - ) - ) - await conn.execute( - text( - "CREATE INDEX IF NOT EXISTS " - "ix_conversations_platform_created_at_inner_id " - "ON conversations (platform_id, created_at DESC, inner_conversation_id DESC)" - ) - ) + await self.conversation_store.migrate() + self.inited = True async def _ensure_persona_folder_columns(self, conn) -> None: """确保 personas 表有 folder_id 和 sort_order 列。 @@ -167,29 +202,60 @@ async def _ensure_persona_custom_error_message_column(self, conn) -> None: ) async def _ensure_platform_message_history_checkpoint_column(self, conn) -> None: - """Ensure platform_message_history has llm_checkpoint_id.""" - result = await conn.execute(text("PRAGMA table_info(platform_message_history)")) - columns = {row[1] for row in result.fetchall()} + """Migrate WebChat associations from legacy markers to event identities. - if "llm_checkpoint_id" not in columns: + Args: + conn: Initialization transaction. + """ + for table, old, new in ( + ("platform_message_history", "llm_checkpoint_id", "turn_id"), + ("webchat_threads", "base_checkpoint_id", "base_event_id"), + ): + columns = { + row[1] + for row in ( + await conn.execute(text(f"PRAGMA table_info({table})")) + ).all() + } + if old in columns and new not in columns: + await conn.execute( + text(f"ALTER TABLE {table} RENAME COLUMN {old} TO {new}") + ) + elif new not in columns: + await conn.execute( + text(f"ALTER TABLE {table} ADD COLUMN {new} VARCHAR") + ) + columns = { + row[1] + for row in ( + await conn.execute(text("PRAGMA table_info(platform_message_history)")) + ).all() + } + if "is_active" not in columns: + await conn.execute( + text( + "ALTER TABLE platform_message_history ADD COLUMN is_active BOOLEAN NOT NULL DEFAULT 1" + ) + ) + if "context_event_id" not in columns: await conn.execute( text( - "ALTER TABLE platform_message_history " - "ADD COLUMN llm_checkpoint_id VARCHAR DEFAULT NULL" + "ALTER TABLE platform_message_history ADD COLUMN context_event_id VARCHAR" ) ) await conn.execute( text( - "CREATE INDEX IF NOT EXISTS " - "ix_platform_message_history_llm_checkpoint_id " - "ON platform_message_history (llm_checkpoint_id)" + "CREATE INDEX IF NOT EXISTS ix_platform_message_history_turn_id ON platform_message_history (turn_id)" ) ) await conn.execute( text( - "CREATE INDEX IF NOT EXISTS " - "ix_platform_message_history_platform_user_id " - "ON platform_message_history (platform_id, user_id, id)" + "CREATE INDEX IF NOT EXISTS ix_platform_message_history_context_event_id ON platform_message_history (context_event_id)" + ) + ) + await conn.execute( + text( + "CREATE INDEX IF NOT EXISTS ix_platform_message_history_platform_user_id ON platform_message_history (platform_id, user_id, id)" ) ) @@ -339,36 +405,47 @@ async def insert_provider_stat( async def get_conversations(self, user_id=None, platform_id=None): async with self.get_db() as session: session: AsyncSession - query = select(ConversationV2) + query = select(ConversationV3) if user_id: - query = query.where(ConversationV2.user_id == user_id) + query = query.where(ConversationV3.umo == user_id) if platform_id: - query = query.where(ConversationV2.platform_id == platform_id) + query = query.where(ConversationV3.platform_id == platform_id) # order by - query = query.order_by(desc(ConversationV2.created_at)) + query = query.order_by(desc(ConversationV3.created_at)) result = await session.execute(query) - return result.scalars().all() + return [ + await self.conversation_store.read_result(session, conv) + for conv in result.scalars() + ] async def get_conversation_by_id(self, cid): async with self.get_db() as session: session: AsyncSession - query = select(ConversationV2).where(ConversationV2.conversation_id == cid) + query = select(ConversationV3).where(ConversationV3.conversation_id == cid) result = await session.execute(query) - return result.scalar_one_or_none() + conv = result.scalar_one_or_none() + return ( + await self.conversation_store.read_result(session, conv) + if conv + else None + ) async def get_all_conversations(self, page=1, page_size=20): async with self.get_db() as session: session: AsyncSession offset = (page - 1) * page_size result = await session.execute( - select(ConversationV2) - .order_by(desc(ConversationV2.created_at)) + select(ConversationV3) + .order_by(desc(ConversationV3.created_at)) .offset(offset) .limit(page_size), ) - return result.scalars().all() + return [ + await self.conversation_store.read_result(session, conv) + for conv in result.scalars() + ] async def get_filtered_conversations( self, @@ -382,39 +459,29 @@ async def get_filtered_conversations( async with self.get_db() as session: session: AsyncSession # Build the base query with filters - base_query = select(ConversationV2) + base_query = select(ConversationV3) conditions = [] if platform_ids: - conditions.append(col(ConversationV2.platform_id).in_(platform_ids)) + conditions.append(col(ConversationV3.platform_id).in_(platform_ids)) # WebChat titles live on the platform session, not on the # conversation row, so the search also matches session titles. if search_query: - escaped_search_query = json.dumps( - search_query, - ensure_ascii=True, - )[1:-1] conditions.append( or_( - col(ConversationV2.title).ilike(f"%{search_query}%"), - col(ConversationV2.user_id).ilike(f"%{search_query}%"), - col(ConversationV2.conversation_id).ilike(f"%{search_query}%"), - col(ConversationV2.content).ilike(f"%{search_query}%"), - col(ConversationV2.content).ilike(f"%{escaped_search_query}%"), + col(ConversationV3.title).ilike(f"%{search_query}%"), + col(ConversationV3.umo).ilike(f"%{search_query}%"), + col(ConversationV3.conversation_id).ilike(f"%{search_query}%"), + _conversation_content_match(search_query), _webchat_session_title_match(search_query), ) ) keyword_query = str(kwargs.get("keyword_query") or "").strip() if keyword_query: - escaped_keyword_query = json.dumps( - keyword_query, - ensure_ascii=True, - )[1:-1] conditions.append( or_( - col(ConversationV2.title).ilike(f"%{keyword_query}%"), - col(ConversationV2.content).ilike(f"%{keyword_query}%"), - col(ConversationV2.content).ilike(f"%{escaped_keyword_query}%"), + col(ConversationV3.title).ilike(f"%{keyword_query}%"), + _conversation_content_match(keyword_query), _webchat_session_title_match(keyword_query), ) ) @@ -423,27 +490,25 @@ async def get_filtered_conversations( conditions.append( or_( *( - col(ConversationV2.user_id).like(f"%:{msg_type}:%") + col(ConversationV3.umo).like(f"%:{msg_type}:%") for msg_type in message_types ) ) ) platforms = kwargs.get("platforms") or [] if platforms: - conditions.append(col(ConversationV2.platform_id).in_(platforms)) + conditions.append(col(ConversationV3.platform_id).in_(platforms)) exclude_ids = kwargs.get("exclude_ids") or [] for exclude_id in exclude_ids: - conditions.append( - not_(col(ConversationV2.user_id).like(f"{exclude_id}%")) - ) + conditions.append(not_(col(ConversationV3.umo).like(f"{exclude_id}%"))) exclude_platforms = kwargs.get("exclude_platforms") or [] if exclude_platforms: conditions.append( - not_(col(ConversationV2.platform_id).in_(exclude_platforms)) + not_(col(ConversationV3.platform_id).in_(exclude_platforms)) ) umo_query = str(kwargs.get("umo_query") or "").strip() if umo_query: - conditions.append(col(ConversationV2.user_id).ilike(f"%{umo_query}%")) + conditions.append(col(ConversationV3.umo).ilike(f"%{umo_query}%")) if conditions: base_query = base_query.where(*conditions) @@ -452,9 +517,9 @@ async def get_filtered_conversations( # Get total count matching the filters count_target = ( - func.distinct(ConversationV2.user_id) + func.distinct(ConversationV3.umo) if group_by_session - else ConversationV2.inner_conversation_id + else ConversationV3.id ) count_query = select(func.count(count_target)) if conditions: @@ -467,23 +532,21 @@ async def get_filtered_conversations( sort_by = kwargs.get("sort_by", "created_at") sort_order = kwargs.get("sort_order", "desc") sort_column = ( - ConversationV2.updated_at + ConversationV3.updated_at if sort_by == "updated_at" - else ConversationV2.created_at + else ConversationV3.created_at ) order = sort_column.asc if sort_order == "asc" else sort_column.desc tie_breaker = ( - ConversationV2.inner_conversation_id.asc - if sort_order == "asc" - else ConversationV2.inner_conversation_id.desc + ConversationV3.id.asc if sort_order == "asc" else ConversationV3.id.desc ) if group_by_session: session_sort = func.max(sort_column).label("session_sort") - session_tie_breaker = func.max( - ConversationV2.inner_conversation_id - ).label("session_tie_breaker") + session_tie_breaker = func.max(ConversationV3.id).label( + "session_tie_breaker" + ) session_query = select( - ConversationV2.user_id, + ConversationV3.umo, session_sort, session_tie_breaker, ) @@ -498,7 +561,7 @@ async def get_filtered_conversations( else session_tie_breaker.desc ) session_rows = await session.execute( - session_query.group_by(ConversationV2.user_id) + session_query.group_by(ConversationV3.umo) .order_by(session_order()) .order_by(session_tie_order()) .offset(offset) @@ -509,11 +572,11 @@ async def get_filtered_conversations( return [], total session_rank = case( {session_id: index for index, session_id in enumerate(session_ids)}, - value=ConversationV2.user_id, + value=ConversationV3.umo, else_=len(session_ids), ) result_query = ( - base_query.where(col(ConversationV2.user_id).in_(session_ids)) + base_query.where(col(ConversationV3.umo).in_(session_ids)) .order_by(session_rank) .order_by(order()) .order_by(tie_breaker()) @@ -525,42 +588,32 @@ async def get_filtered_conversations( .offset(offset) .limit(page_size) ) - if not include_history: - result_query = result_query.options(defer(ConversationV2.content)) if ( not group_by_session and sort_by == "created_at" and (len(platforms) > 1 or len(platform_ids or []) > 1) ): - # SQLite may choose the narrow platform index for IN queries and - # then materialize a temporary sort. Force the global ordering - # index for multi-platform pages while keeping ORM row mapping. compiled = result_query.compile( dialect=sqlite_dialect(paramstyle="named"), compile_kwargs={"render_postcompile": True}, ) indexed_sql = compiled.string.replace( - "FROM conversations", - "FROM conversations INDEXED BY " - "ix_conversations_created_at_inner_id", + "FROM conversations_v3", + "FROM conversations_v3 INDEXED BY ix_conversations_v3_created_id", 1, ) - conversation_columns = [ - column - for column in ConversationV2.__table__.columns - if include_history or column.name != "content" - ] - result_query = select(ConversationV2).from_statement( - text(indexed_sql).columns(*conversation_columns), - ) - if not include_history: - result_query = result_query.options( - defer(ConversationV2.content), - ) + result_query = select(ConversationV3).from_statement( + text(indexed_sql).columns(*ConversationV3.__table__.columns) + ) result = await session.execute(result_query, compiled.params) else: result = await session.execute(result_query) - conversations = result.scalars().all() + conversations = [ + await self.conversation_store.read_result( + session, conv, include_history + ) + for conv in result.scalars() + ] return conversations, total @@ -572,9 +625,9 @@ async def get_conversation_platform_ids(self) -> list[str]: """ async with self.get_db() as session: result = await session.execute( - select(ConversationV2.platform_id) + select(ConversationV3.platform_id) .distinct() - .order_by(ConversationV2.platform_id) + .order_by(ConversationV3.platform_id) ) return [platform_id for platform_id in result.scalars() if platform_id] @@ -589,70 +642,50 @@ async def create_conversation( created_at=None, updated_at=None, ): - kwargs = {} - if cid: - kwargs["conversation_id"] = cid - if created_at: - kwargs["created_at"] = created_at - if updated_at: - kwargs["updated_at"] = updated_at + conv = await self.conversation_store.create( + umo=user_id, + platform_id=platform_id, + content=content, + title=title, + persona_id=persona_id, + cid=cid, + created_at=created_at, + updated_at=updated_at, + ) async with self.get_db() as session: - session: AsyncSession - async with session.begin(): - new_conversation = ConversationV2( - user_id=user_id, - content=content or [], - platform_id=platform_id, - title=title, - persona_id=persona_id, - **kwargs, - ) - session.add(new_conversation) - return new_conversation + return await self.conversation_store.read_result(session, conv) async def update_conversation( self, cid, title=None, persona_id=None, content=None, token_usage=None ): - async with self.get_db() as session: - session: AsyncSession - async with session.begin(): - query = update(ConversationV2).where( - col(ConversationV2.conversation_id) == cid, - ) - values = {} - if title is not None: - values["title"] = title - if persona_id is not None: - values["persona_id"] = persona_id - if content is not None: - values["content"] = content - if token_usage is not None: - values["token_usage"] = token_usage - if not values: - return None - query = query.values(**values) - await session.execute(query) + snapshot = await self.conversation_store.read(cid) + if snapshot is None: + return None + changes = {} + if title is not None: + changes["title"] = title + if persona_id is not None: + changes["persona_id"] = persona_id + drafts = [] + if changes or token_usage is not None: + payload = {"changes": changes} + if token_usage is not None: + payload["token_usage"] = token_usage + drafts.append({"type": "conversation.updated", "payload": payload}) + await self.conversation_store.append( + cid, + drafts, + history=content, + expected_head=snapshot.conversation.head_seq, + expected_leaf=snapshot.conversation.leaf_event_id, + ) return await self.get_conversation_by_id(cid) async def delete_conversation(self, cid) -> None: - async with self.get_db() as session: - session: AsyncSession - async with session.begin(): - await session.execute( - delete(ConversationV2).where( - col(ConversationV2.conversation_id) == cid, - ), - ) + await self.conversation_store.delete(cid=cid) async def delete_conversations_by_user_id(self, user_id: str) -> None: - async with self.get_db() as session: - session: AsyncSession - async with session.begin(): - await session.execute( - delete(ConversationV2).where( - col(ConversationV2.user_id) == user_id - ), - ) + await self.conversation_store.delete(umo=user_id) async def get_session_conversations( self, @@ -672,19 +705,22 @@ async def get_session_conversations( func.json_extract(Preference.value, "$.val").label( "conversation_id", ), # type: ignore - col(ConversationV2.persona_id).label("persona_id"), - col(ConversationV2.title).label("title"), + col(ConversationV3.persona_id).label("persona_id"), + col(ConversationV3.title).label("title"), col(Persona.persona_id).label("persona_name"), ) .select_from(Preference) .outerjoin( - ConversationV2, + ConversationV3, + ConversationEvent, + ConversationV3, + ConversationEvent, func.json_extract(Preference.value, "$.val") - == ConversationV2.conversation_id, + == ConversationV3.conversation_id, ) .outerjoin( Persona, - col(ConversationV2.persona_id) == Persona.persona_id, + col(ConversationV3.persona_id) == Persona.persona_id, ) .where(Preference.scope == "umo", Preference.key == "sel_conv_id") ) @@ -695,7 +731,7 @@ async def get_session_conversations( base_query = base_query.where( or_( col(Preference.scope_id).ilike(search_pattern), - col(ConversationV2.title).ilike(search_pattern), + col(ConversationV3.title).ilike(search_pattern), col(Persona.persona_id).ilike(search_pattern), ), ) @@ -720,13 +756,16 @@ async def get_session_conversations( select(func.count(col(Preference.scope_id))) .select_from(Preference) .outerjoin( - ConversationV2, + ConversationV3, + ConversationEvent, + ConversationV3, + ConversationEvent, func.json_extract(Preference.value, "$.val") - == ConversationV2.conversation_id, + == ConversationV3.conversation_id, ) .outerjoin( Persona, - col(ConversationV2.persona_id) == Persona.persona_id, + col(ConversationV3.persona_id) == Persona.persona_id, ) .where(Preference.scope == "umo", Preference.key == "sel_conv_id") ) @@ -737,7 +776,7 @@ async def get_session_conversations( count_base_query = count_base_query.where( or_( col(Preference.scope_id).ilike(search_pattern), - col(ConversationV2.title).ilike(search_pattern), + col(ConversationV3.title).ilike(search_pattern), col(Persona.persona_id).ilike(search_pattern), ), ) @@ -770,21 +809,75 @@ async def insert_platform_message_history( content, sender_id=None, sender_name=None, - llm_checkpoint_id=None, + turn_id=None, max_messages=None, + llm_checkpoint_id: str | None = None, ): """Insert a new platform message history record.""" + if turn_id is None: + turn_id = llm_checkpoint_id async with self.get_db() as session: session: AsyncSession async with session.begin(): + if turn_id: + # Serialize link lookup and insertion with context-event writes. + # Otherwise both writers can miss the other one's pending row. + await session.execute(text("BEGIN IMMEDIATE")) new_history = PlatformMessageHistory( platform_id=platform_id, user_id=user_id, content=content, sender_id=sender_id, sender_name=sender_name, - llm_checkpoint_id=llm_checkpoint_id, + turn_id=turn_id, ) + if turn_id: + role = "user" if content.get("type") == "user" else "assistant" + turn_conversation = ( + select(ConversationEvent.conversation_ref) + .where( + ConversationEvent.event_id == turn_id, + ConversationEvent.type == "turn.started", + ) + .scalar_subquery() + ) + query = ( + select(ConversationEvent.event_id) + .where( + ConversationEvent.conversation_ref == turn_conversation, + ConversationEvent.seq + > select(ConversationEvent.seq) + .where(ConversationEvent.event_id == turn_id) + .scalar_subquery(), + ConversationEvent.payload["turn_id"].as_string() == turn_id, + or_( + (ConversationEvent.type == "message.appended") + & ( + ConversationEvent.payload["message"][ + "role" + ].as_string() + == role + ), + (ConversationEvent.type == "context.rebased") + & ( + func.json_extract( + ConversationEvent.payload, + "$.messages[#-1].message.role", + ) + == role + ), + ), + ) + .order_by( + ConversationEvent.seq.asc() + if role == "user" + else ConversationEvent.seq.desc() + ) + .limit(1) + ) + new_history.context_event_id = ( + await session.execute(query) + ).scalar_one_or_none() session.add(new_history) await session.flush() if max_messages is not None: @@ -810,14 +903,18 @@ async def update_platform_message_history( self, message_id: int, content: dict | None = None, + turn_id: str | None = None, llm_checkpoint_id: str | None = None, ) -> None: """Update a platform message history record.""" + if turn_id is None: + turn_id = llm_checkpoint_id values = {} if content is not None: values["content"] = content - if llm_checkpoint_id is not None: - values["llm_checkpoint_id"] = llm_checkpoint_id + if turn_id is not None: + values["turn_id"] = turn_id + values["context_event_id"] = None if not values: return @@ -877,6 +974,7 @@ async def get_platform_message_history( .where( PlatformMessageHistory.platform_id == platform_id, PlatformMessageHistory.user_id == user_id, + PlatformMessageHistory.is_active == True, # noqa: E712 ) .order_by( desc(PlatformMessageHistory.created_at), @@ -927,7 +1025,7 @@ async def create_webchat_thread( creator: str, parent_session_id: str, parent_message_id: int, - base_checkpoint_id: str, + base_event_id: str, selected_text: str, ) -> WebChatThread: """Create a WebChat side thread.""" @@ -938,7 +1036,7 @@ async def create_webchat_thread( creator=creator, parent_session_id=parent_session_id, parent_message_id=parent_message_id, - base_checkpoint_id=base_checkpoint_id, + base_event_id=base_event_id, selected_text=selected_text, ) session.add(thread) diff --git a/astrbot/core/pipeline/context_utils.py b/astrbot/core/pipeline/context_utils.py index 5497bbf974..f00cb8be4b 100644 --- a/astrbot/core/pipeline/context_utils.py +++ b/astrbot/core/pipeline/context_utils.py @@ -3,6 +3,10 @@ import typing as T from astrbot import logger +from astrbot.core.agent.conversation_events import ( + active_conversation_writer, + active_plugin_id, +) from astrbot.core.message.message_event_result import CommandResult, MessageEventResult from astrbot.core.platform.astr_message_event import AstrMessageEvent from astrbot.core.star.star import star_map @@ -98,7 +102,15 @@ async def call_event_hook( logger.debug( f"hook({hook_type.name}) -> {star_map[handler.handler_module_path].name} - {handler.handler_name}", ) - await handler.handler(event, *args, **kwargs) + writer_token = active_conversation_writer.set(event.conversation_events) + plugin_token = active_plugin_id.set( + star_map[handler.handler_module_path].name + ) + try: + await handler.handler(event, *args, **kwargs) + finally: + active_plugin_id.reset(plugin_token) + active_conversation_writer.reset(writer_token) except BaseException: logger.error(traceback.format_exc()) diff --git a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py index 2f7a34b2b3..64f87d2272 100644 --- a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py +++ b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/internal.py @@ -6,9 +6,8 @@ from dataclasses import replace from astrbot.core import db_helper, logger +from astrbot.core.agent.conversation_events import ConversationEventWriter from astrbot.core.agent.message import ( - CheckpointData, - CheckpointMessageSegment, Message, dump_messages_with_checkpoints, ) @@ -206,6 +205,7 @@ async def process( logger.debug("acquired session lock for llm request") agent_runner: AgentRunner | None = None runner_registered = False + event_writer = None try: build_cfg = replace( self.main_agent_cfg, @@ -232,6 +232,7 @@ async def process( agent_runner = build_result.agent_runner req = build_result.provider_request + event_writer = build_result.conversation_events provider = build_result.provider reset_coro = build_result.reset_coro @@ -306,6 +307,7 @@ async def process( self.show_tool_call_result, show_reasoning=show_reasoning, buffer_intermediate_messages=self.buffer_intermediate_messages, + conversation_events=event_writer, ), ), ) @@ -322,6 +324,7 @@ async def process( agent_runner.run_context.messages, agent_runner.stats, user_aborted=agent_runner.was_aborted(), + conversation_events=event_writer, ) elif streaming_response and not stream_to_general: @@ -337,6 +340,7 @@ async def process( self.show_tool_call_result, show_reasoning=show_reasoning, buffer_intermediate_messages=self.buffer_intermediate_messages, + conversation_events=event_writer, ), ), ) @@ -368,6 +372,7 @@ async def process( stream_to_general, show_reasoning=show_reasoning, buffer_intermediate_messages=self.buffer_intermediate_messages, + conversation_events=event_writer, ): yield @@ -397,6 +402,7 @@ async def process( agent_runner.run_context.messages, agent_runner.stats, user_aborted=agent_runner.was_aborted(), + conversation_events=event_writer, ) asyncio.create_task( @@ -407,8 +413,27 @@ async def process( ), ) finally: - if runner_registered and agent_runner is not None: - unregister_active_runner(event.unified_msg_origin, agent_runner) + try: + if event_writer and event_writer.turn_id: + status = "failed" + if agent_runner and agent_runner.was_aborted(): + status = "cancelled" + elif agent_runner and agent_runner.done(): + final = agent_runner.get_final_llm_resp() + status = ( + "completed" + if final and final.role == "assistant" + else "failed" + ) + if event.is_stopped(): + status = "cancelled" + await event_writer.finish_turn(status) + finally: + event.conversation_events = None + if runner_registered and agent_runner is not None: + unregister_active_runner( + event.unified_msg_origin, agent_runner + ) except Exception as e: logger.error(f"Error occurred while processing agent: {e}") @@ -440,40 +465,47 @@ async def _save_to_history( all_messages: list[Message], runner_stats: AgentStats | None, user_aborted: bool = False, + *, + conversation_events: ConversationEventWriter | None = None, ) -> None: if not req or not req.conversation: return + writer = conversation_events messages_to_save: list[Message] = [] - skipped_initial_system = False - for message in all_messages: - if message.role == "system" and not skipped_initial_system: - skipped_initial_system = True + for index, message in enumerate(all_messages): + if index == 0 and message.role == "system": continue - if message.role in ["assistant", "user"] and message._no_save: + if ( + not writer + and message.role in ["assistant", "user"] + and message._no_save + ): continue messages_to_save.append(message) - checkpoint_id = event.get_extra("llm_checkpoint_id") - has_checkpoint = isinstance(checkpoint_id, str) and bool(checkpoint_id) - message_to_save = dump_messages_with_checkpoints(messages_to_save) + turn_id = event.get_extra("turn_id") + has_turn = isinstance(turn_id, str) and bool(turn_id) + message_to_save = [ + m + for m in dump_messages_with_checkpoints( + messages_to_save, include_temporary=bool(writer) + ) + if m.get("role") != "_checkpoint" + ] if not user_aborted and ( llm_response is None or llm_response.role != "assistant" ): - if has_checkpoint: - message_to_save.append( - CheckpointMessageSegment( - content=CheckpointData(id=checkpoint_id), - ).model_dump() - ) - if has_checkpoint or (llm_response is None and req.tool_calls_result): - token_usage = None if has_checkpoint else req.conversation.token_usage - await self.conv_manager.update_conversation( - event.unified_msg_origin, - req.conversation.cid, - history=message_to_save, - token_usage=token_usage, - ) + if has_turn or (llm_response is None and req.tool_calls_result): + if writer: + await writer.save_history(message_to_save) + else: + await self.conv_manager.update_conversation( + event.unified_msg_origin, + req.conversation.cid, + history=message_to_save, + token_usage=None if has_turn else req.conversation.token_usage, + ) return if llm_response and llm_response.role != "assistant": @@ -494,32 +526,20 @@ async def _save_to_history( logger.debug("The LLM response is empty; not saving a record.") return - if isinstance(checkpoint_id, str) and checkpoint_id: - message_to_save.append( - CheckpointMessageSegment( - content=CheckpointData(id=checkpoint_id), - ).model_dump() - ) - - # if user_aborted: - # message_to_save.append( - # Message( - # role="assistant", - # content="[User aborted this request. Partial output before abort was preserved.]", - # ).model_dump() - # ) - token_usage = None if runner_stats: # token_usage = runner_stats.token_usage.total token_usage = llm_response.usage.total if llm_response.usage else None - await self.conv_manager.update_conversation( - event.unified_msg_origin, - req.conversation.cid, - history=message_to_save, - token_usage=token_usage, - ) + if writer: + await writer.save_history(message_to_save, token_usage=token_usage) + else: + await self.conv_manager.update_conversation( + event.unified_msg_origin, + req.conversation.cid, + history=message_to_save, + token_usage=token_usage, + ) # we prevent astrbot from connecting to known malicious hosts diff --git a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py index 93836c1fe8..2a61db4149 100644 --- a/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py +++ b/astrbot/core/pipeline/process_stage/method/agent_sub_stages/third_party.py @@ -14,6 +14,7 @@ ) from astrbot.core.agent.runners.dify.dify_agent_runner import DifyAgentRunner from astrbot.core.astr_agent_hooks import MAIN_AGENT_HOOKS +from astrbot.core.conversation_mgr import ConversationManager from astrbot.core.message.components import Image, Record from astrbot.core.message.message_event_result import ( MessageChain, @@ -303,6 +304,18 @@ async def process( custom_error_message = await self._resolve_persona_custom_error_message(event) set_persona_custom_error_message_on_event(event, custom_error_message) + event_writer = None + manager = self.ctx.plugin_manager.context.conversation_manager + if isinstance(manager, ConversationManager): + cid = await manager.get_curr_conversation_id(event.unified_msg_origin) + if not cid: + cid = await manager.new_conversation( + event.unified_msg_origin, event.get_platform_id() + ) + event_writer = await manager.event_writer(event.unified_msg_origin, cid) + event_writer.request = req + event.conversation_events = event_writer + # call event hook if await call_event_hook(event, EventType.OnLLMRequestEvent, req): return @@ -344,7 +357,26 @@ async def close_runner_once() -> None: if runner_closed: return runner_closed = True - await _close_runner_if_supported(runner) + try: + if event_writer and event_writer.turn_id: + final = runner.get_final_llm_resp() + status = ( + "completed" if final and final.role == "assistant" else "failed" + ) + if event.is_stopped() or not runner.done(): + status = "cancelled" + if status == "completed": + history = [entry["message"] for entry in event_writer._entries] + history.append(await req.assemble_context()) + if final.completion_text: + history.append( + {"role": "assistant", "content": final.completion_text} + ) + await event_writer.save_history(history) + await event_writer.finish_turn(status) + finally: + event.conversation_events = None + await _close_runner_if_supported(runner) def mark_stream_consumed() -> None: nonlocal stream_consumed @@ -353,6 +385,17 @@ def mark_stream_consumed() -> None: stream_watchdog_task.cancel() try: + if event_writer: + turn_id = event.get_extra("turn_id") + await event_writer.start_turn( + { + "kind": "im_wake", + "umo": event.unified_msg_origin, + "runner": self.runner_type, + "context_mode": "transcript_only", + }, + event_id=turn_id if isinstance(turn_id, str) else None, + ) await runner.reset( request=req, run_context=AgentContextWrapper( diff --git a/astrbot/core/platform/astr_message_event.py b/astrbot/core/platform/astr_message_event.py index c9b5bd67df..50f257ef0d 100644 --- a/astrbot/core/platform/astr_message_event.py +++ b/astrbot/core/platform/astr_message_event.py @@ -6,7 +6,7 @@ import uuid from collections.abc import AsyncGenerator from time import time -from typing import Any +from typing import TYPE_CHECKING, Any from deprecated import deprecated @@ -33,6 +33,9 @@ from .message_session import MessageSesion, MessageSession # noqa from .platform_metadata import PlatformMetadata +if TYPE_CHECKING: + from astrbot.core.agent.conversation_events import ConversationEventWriter + class AstrMessageEvent(abc.ABC): def __init__( @@ -56,6 +59,8 @@ def __init__( """是否是 At 机器人或者带有唤醒词或者是私聊(插件注册的事件监听器会让 is_wake 设为 True, 但是不会让这个属性置为 True)""" self._extras: dict[str, Any] = {} self._force_stopped: bool = False + self.conversation_events: ConversationEventWriter | None = None + """Host-bound conversation writer available to hooks during an agent run.""" """独立的停止标志,不依赖 _result,不会被 clear_result() 重置""" message_type = getattr(message_obj, "type", None) if not isinstance(message_type, MessageType): diff --git a/astrbot/core/platform/sources/webchat/webchat_adapter.py b/astrbot/core/platform/sources/webchat/webchat_adapter.py index 535b214a3e..84ecee5682 100644 --- a/astrbot/core/platform/sources/webchat/webchat_adapter.py +++ b/astrbot/core/platform/sources/webchat/webchat_adapter.py @@ -265,9 +265,7 @@ def create_event(self, message: AstrBotMessage) -> WebChatMessageEvent: ) message_event.set_extra("selected_model", payload.get("selected_model")) message_event.set_extra("action_type", payload.get("action_type")) - message_event.set_extra( - "llm_checkpoint_id", payload.get("llm_checkpoint_id") - ) + message_event.set_extra("turn_id", payload.get("turn_id")) message_event.set_extra( "thread_selected_text", payload.get("thread_selected_text") ) diff --git a/astrbot/core/platform_message_history_mgr.py b/astrbot/core/platform_message_history_mgr.py index 12dc2fff19..4a9c9543ea 100644 --- a/astrbot/core/platform_message_history_mgr.py +++ b/astrbot/core/platform_message_history_mgr.py @@ -15,17 +15,20 @@ async def insert( content: dict, sender_id: str | None = None, sender_name: str | None = None, - llm_checkpoint_id: str | None = None, + turn_id: str | None = None, max_messages: int | None = None, + llm_checkpoint_id: str | None = None, ) -> PlatformMessageHistory: """Insert a new platform message history record.""" + if turn_id is None: + turn_id = llm_checkpoint_id return await self.db.insert_platform_message_history( platform_id=platform_id, user_id=user_id, content=content, sender_id=sender_id, sender_name=sender_name, - llm_checkpoint_id=llm_checkpoint_id, + turn_id=turn_id, max_messages=max_messages, ) @@ -151,13 +154,16 @@ async def update( self, message_id: int, content: dict | None = None, + turn_id: str | None = None, llm_checkpoint_id: str | None = None, ) -> None: """Update a platform message history record.""" + if turn_id is None: + turn_id = llm_checkpoint_id await self.db.update_platform_message_history( message_id=message_id, content=content, - llm_checkpoint_id=llm_checkpoint_id, + turn_id=turn_id, ) async def delete_by_id(self, message_id: int) -> None: diff --git a/astrbot/core/provider/sources/anthropic_source.py b/astrbot/core/provider/sources/anthropic_source.py index c861ded6ba..809888808c 100644 --- a/astrbot/core/provider/sources/anthropic_source.py +++ b/astrbot/core/provider/sources/anthropic_source.py @@ -12,6 +12,10 @@ from astrbot import logger from astrbot.api.provider import Provider +from astrbot.core.agent.event_stream import ( + RequestEventRecorder, + request_recorder_kwargs, +) from astrbot.core.agent.message import AudioURLPart, ContentPart, ImageURLPart, TextPart from astrbot.core.exceptions import EmptyModelOutputError from astrbot.core.provider.entities import LLMResponse, TokenUsage @@ -512,6 +516,7 @@ async def _query( tools: ToolSet | None, *, request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, ) -> LLMResponse: if tools: if tool_list := tools.get_func_desc_anthropic_style(): @@ -535,6 +540,7 @@ async def _query( **payloads, stream=False, extra_body=extra_body ), max_attempts=request_max_retries, + request_event_recorder=request_event_recorder, ) except httpx.RequestError as e: proxy = self.provider_config.get("proxy", "") @@ -605,6 +611,7 @@ async def _query_stream( tools: ToolSet | None, *, request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, ) -> AsyncGenerator[LLMResponse, None]: if tools: if tool_list := tools.get_func_desc_anthropic_style(): @@ -634,6 +641,7 @@ async def _query_stream( "Anthropic", lambda: self.client.messages.stream(**payloads, extra_body=extra_body), max_attempts=request_max_retries, + request_event_recorder=request_event_recorder, ) as stream: assert isinstance(stream, anthropic.AsyncMessageStream) async for event in stream: @@ -773,6 +781,7 @@ async def text_chat( extra_user_content_parts=None, tool_choice: Literal["auto", "any", "tool", "none"] | dict[str, str] = "auto", request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, **kwargs, ) -> LLMResponse: if contexts is None: @@ -826,6 +835,7 @@ async def text_chat( payloads, func_tool, request_max_retries=request_max_retries, + **request_recorder_kwargs(self._query, request_event_recorder), ) except Exception as e: raise e @@ -846,6 +856,7 @@ async def text_chat_stream( extra_user_content_parts=None, tool_choice: Literal["auto", "any", "tool", "none"] | dict[str, str] = "auto", request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, **kwargs, ): if contexts is None: @@ -896,6 +907,7 @@ async def text_chat_stream( payloads, func_tool, request_max_retries=request_max_retries, + **request_recorder_kwargs(self._query_stream, request_event_recorder), ): yield llm_response diff --git a/astrbot/core/provider/sources/gemini_source.py b/astrbot/core/provider/sources/gemini_source.py index b6b7a97fb2..9028cb820d 100644 --- a/astrbot/core/provider/sources/gemini_source.py +++ b/astrbot/core/provider/sources/gemini_source.py @@ -14,6 +14,10 @@ import astrbot.core.message.components as Comp from astrbot import logger from astrbot.api.provider import Provider +from astrbot.core.agent.event_stream import ( + RequestEventRecorder, + request_recorder_kwargs, +) from astrbot.core.agent.message import AudioURLPart, ContentPart, ImageURLPart, TextPart from astrbot.core.exceptions import EmptyModelOutputError from astrbot.core.message.message_event_result import MessageChain @@ -602,6 +606,7 @@ async def _query( tools: ToolSet | None, *, request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, ) -> LLMResponse: """非流式请求 Gemini API""" system_instruction = next( @@ -637,6 +642,7 @@ async def _query( config=config, ), max_attempts=request_max_retries, + request_event_recorder=request_event_recorder, ) logger.debug(f"genai result: {result}") @@ -703,6 +709,7 @@ async def _query_stream( tools: ToolSet | None, *, request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, ) -> AsyncGenerator[LLMResponse, None]: """流式请求 Gemini API""" system_instruction = next( @@ -729,6 +736,7 @@ async def _query_stream( config=config, ), max_attempts=request_max_retries, + request_event_recorder=request_event_recorder, ) break except APIError as e: @@ -845,6 +853,7 @@ async def text_chat( extra_user_content_parts=None, tool_choice: Literal["auto", "required"] = "auto", request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, **kwargs, ) -> LLMResponse: if contexts is None: @@ -890,6 +899,7 @@ async def text_chat( payloads, func_tool, request_max_retries=request_max_retries, + **request_recorder_kwargs(self._query, request_event_recorder), ) except APIError as e: if await self._handle_api_error(e, keys): @@ -912,6 +922,7 @@ async def text_chat_stream( extra_user_content_parts=None, tool_choice: Literal["auto", "required"] = "auto", request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, **kwargs, ) -> AsyncGenerator[LLMResponse, None]: if contexts is None: @@ -957,6 +968,9 @@ async def text_chat_stream( payloads, func_tool, request_max_retries=request_max_retries, + **request_recorder_kwargs( + self._query_stream, request_event_recorder + ), ): yield response break diff --git a/astrbot/core/provider/sources/openai_responses_source.py b/astrbot/core/provider/sources/openai_responses_source.py index c5cb9bdb82..00c2334774 100644 --- a/astrbot/core/provider/sources/openai_responses_source.py +++ b/astrbot/core/provider/sources/openai_responses_source.py @@ -8,6 +8,7 @@ import astrbot.core.message.components as Comp from astrbot import logger +from astrbot.core.agent.event_stream import RequestEventRecorder from astrbot.core.agent.message import ContentPart, Message from astrbot.core.agent.tool import ToolSet from astrbot.core.exceptions import EmptyModelOutputError @@ -300,6 +301,7 @@ async def _query( tools: ToolSet | None, *, request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, ) -> LLMResponse: """Send a non-streaming Responses API request. @@ -353,6 +355,7 @@ async def _query( extra_body=extra_body, ), max_attempts=request_max_retries, + request_event_recorder=request_event_recorder, ) if not isinstance(response, Response): raise TypeError( @@ -369,6 +372,7 @@ async def _query_stream( tools: ToolSet | None, *, request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, ) -> AsyncGenerator[LLMResponse, None]: """Send a streaming Responses API request. @@ -422,6 +426,7 @@ async def _query_stream( extra_body=extra_body, ), max_attempts=request_max_retries, + request_event_recorder=request_event_recorder, ) response_id: str | None = None diff --git a/astrbot/core/provider/sources/openai_source.py b/astrbot/core/provider/sources/openai_source.py index f7870b7137..3df39fce5a 100644 --- a/astrbot/core/provider/sources/openai_source.py +++ b/astrbot/core/provider/sources/openai_source.py @@ -18,6 +18,10 @@ import astrbot.core.message.components as Comp from astrbot import logger from astrbot.api.provider import Provider +from astrbot.core.agent.event_stream import ( + RequestEventRecorder, + request_recorder_kwargs, +) from astrbot.core.agent.message import ( AudioURLPart, ContentPart, @@ -534,6 +538,7 @@ async def _query( tools: ToolSet | None, *, request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, ) -> LLMResponse: if tools: model = payloads.get("model", "").lower() @@ -573,6 +578,7 @@ async def _query( extra_body=extra_body, ), max_attempts=request_max_retries, + request_event_recorder=request_event_recorder, ) if not isinstance(completion, ChatCompletion): @@ -592,6 +598,7 @@ async def _query_stream( tools: ToolSet | None, *, request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, ) -> AsyncGenerator[LLMResponse, None]: """流式查询API,逐步返回结果""" if tools: @@ -632,6 +639,7 @@ async def _query_stream( stream_options={"include_usage": True}, ), max_attempts=request_max_retries, + request_event_recorder=request_event_recorder, ) llm_response = LLMResponse("assistant", is_chunk=True) @@ -1196,6 +1204,7 @@ async def text_chat( extra_user_content_parts=None, tool_choice: Literal["auto", "required"] = "auto", request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, **kwargs, ) -> LLMResponse: payloads, context_query = await self._prepare_chat_payload( @@ -1227,6 +1236,7 @@ async def text_chat( payloads, func_tool, request_max_retries=request_max_retries, + **request_recorder_kwargs(self._query, request_event_recorder), ) break except Exception as e: @@ -1273,6 +1283,7 @@ async def text_chat_stream( model=None, tool_choice: Literal["auto", "required"] = "auto", request_max_retries: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, **kwargs, ) -> AsyncGenerator[LLMResponse, None]: """流式对话,与服务商交互并逐步返回结果""" @@ -1303,6 +1314,9 @@ async def text_chat_stream( payloads, func_tool, request_max_retries=request_max_retries, + **request_recorder_kwargs( + self._query_stream, request_event_recorder + ), ): yield response break diff --git a/astrbot/core/provider/sources/request_retry.py b/astrbot/core/provider/sources/request_retry.py index 14dd57d7ab..d0b95a8172 100644 --- a/astrbot/core/provider/sources/request_retry.py +++ b/astrbot/core/provider/sources/request_retry.py @@ -11,6 +11,7 @@ ) from astrbot import logger +from astrbot.core.agent.event_stream import RequestEventRecorder from astrbot.core.utils.config_number import coerce_int_config from astrbot.core.utils.network_utils import is_connection_error @@ -114,6 +115,7 @@ async def retry_provider_request( *, retry_rate_limits: bool = True, max_attempts: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, ) -> T: retrying = _build_retrying( provider_label, @@ -123,7 +125,15 @@ async def retry_provider_request( async for attempt in retrying: with attempt: - return await request_factory() + recorder = request_event_recorder + if recorder: + await recorder.before_network_attempt() + try: + return await request_factory() + except Exception as exc: + if recorder: + await recorder.finish("failed", error_code=type(exc).__name__) + raise raise RuntimeError("Provider request retry loop exited unexpectedly.") @@ -135,6 +145,7 @@ async def retry_provider_request_context( *, retry_rate_limits: bool = True, max_attempts: int | None = None, + request_event_recorder: RequestEventRecorder | None = None, ) -> AsyncIterator[T]: manager: AbstractAsyncContextManager[T] | None = None @@ -148,6 +159,7 @@ async def _enter_context() -> T: _enter_context, retry_rate_limits=retry_rate_limits, max_attempts=max_attempts, + request_event_recorder=request_event_recorder, ) if manager is None: diff --git a/astrbot/core/utils/history_saver.py b/astrbot/core/utils/history_saver.py index 840d3f1871..02b2962574 100644 --- a/astrbot/core/utils/history_saver.py +++ b/astrbot/core/utils/history_saver.py @@ -1,6 +1,7 @@ import json from astrbot import logger +from astrbot.core.agent.conversation_events import ConversationEventWriter from astrbot.core.conversation_mgr import ConversationManager from astrbot.core.platform.astr_message_event import AstrMessageEvent from astrbot.core.provider.entities import ProviderRequest @@ -12,8 +13,17 @@ async def persist_agent_history( event: AstrMessageEvent, req: ProviderRequest, summary_note: str, + conversation_events: ConversationEventWriter | None = None, ) -> None: - """Persist agent interaction into conversation history.""" + """Persist an autonomous task summary at its host's history-save boundary. + + Args: + conversation_manager: Legacy conversation API used without a bound writer. + event: Event identifying the conversation owner. + req: Request with the original conversation history. + summary_note: Task result added to future context. + conversation_events: Explicit host collector for the running turn. + """ if not req or not req.conversation: return @@ -24,6 +34,9 @@ async def persist_agent_history( logger.warning("Failed to parse conversation history: %s", exc) history.append({"role": "user", "content": "Output your last task result below."}) history.append({"role": "assistant", "content": summary_note}) + if conversation_events is not None: + await conversation_events.save_history(history) + return await conversation_manager.update_conversation( event.unified_msg_origin, req.conversation.cid, diff --git a/astrbot/dashboard/services/chat_service.py b/astrbot/dashboard/services/chat_service.py index ca651f6452..a888a7151e 100644 --- a/astrbot/dashboard/services/chat_service.py +++ b/astrbot/dashboard/services/chat_service.py @@ -12,7 +12,6 @@ from typing import Any from astrbot.core import logger, sp -from astrbot.core.agent.message import get_checkpoint_id, is_checkpoint_message from astrbot.core.core_lifecycle import AstrBotCoreLifecycle from astrbot.core.db import BaseDatabase from astrbot.core.platform.message_type import MessageType @@ -326,7 +325,7 @@ def serialize_thread(thread) -> dict: "thread_id": thread.thread_id, "parent_session_id": thread.parent_session_id, "parent_message_id": thread.parent_message_id, - "base_checkpoint_id": thread.base_checkpoint_id, + "base_event_id": thread.base_event_id, "selected_text": thread.selected_text, "created_at": to_utc_isoformat(thread.created_at), "updated_at": to_utc_isoformat(thread.updated_at), @@ -350,119 +349,6 @@ def serialize_history_entry(history) -> dict: } -def find_checkpoint_index(history: list[dict], checkpoint_id: str) -> int | None: - for index, message in enumerate(history): - if get_checkpoint_id(message) == checkpoint_id: - return index - return None - - -def find_turn_range(history: list[dict], checkpoint_id: str) -> tuple[int, int] | None: - checkpoint_index = find_checkpoint_index(history, checkpoint_id) - if checkpoint_index is None: - return None - - start = 0 - for index in range(checkpoint_index - 1, -1, -1): - if is_checkpoint_message(history[index]): - start = index + 1 - break - return start, checkpoint_index - - -def is_latest_checkpoint(history: list[dict], checkpoint_id: str) -> bool: - for message in reversed(history): - current_checkpoint_id = get_checkpoint_id(message) - if current_checkpoint_id: - return current_checkpoint_id == checkpoint_id - return False - - -def replace_user_conversation_content(original_content, edited_text: str): - if isinstance(original_content, str): - return edited_text - if not isinstance(original_content, list): - return edited_text - - result: list[dict] = [] - inserted_text = False - for part in original_content: - if not isinstance(part, dict): - result.append(part) - continue - if part.get("type") != "text": - result.append(part) - continue - text = part.get("text") - if isinstance(text, str) and text.startswith(""): - result.append(part) - continue - if not inserted_text and edited_text: - result.append({"type": "text", "text": edited_text}) - inserted_text = True - - if not inserted_text and edited_text: - result.insert(0, {"type": "text", "text": edited_text}) - return result - - -def replace_assistant_conversation_content( - original_content, - edited_text: str, - reasoning: str, -): - if isinstance(original_content, str): - return edited_text - if not isinstance(original_content, list): - return [{"type": "text", "text": edited_text}] if edited_text else [] - - result: list[dict] = [] - inserted_text = False - inserted_think = False - for part in original_content: - if not isinstance(part, dict): - result.append(part) - continue - if part.get("type") == "text": - if not inserted_text and edited_text: - result.append({"type": "text", "text": edited_text}) - inserted_text = True - continue - if part.get("type") == "think": - if not inserted_think and reasoning: - result.append({"type": "think", "think": reasoning}) - inserted_think = True - continue - result.append(part) - - if reasoning and not inserted_think: - result.insert(0, {"type": "think", "think": reasoning}) - if edited_text and not inserted_text: - result.append({"type": "text", "text": edited_text}) - return result - - -def find_turn_user_index(history: list[dict], start: int, end: int) -> int | None: - for index in range(start, end): - message = history[index] - if isinstance(message, dict) and message.get("role") == "user": - return index - return None - - -def find_turn_final_assistant_index( - history: list[dict], start: int, end: int -) -> int | None: - for index in range(end - 1, start - 1, -1): - message = history[index] - if not isinstance(message, dict) or message.get("role") != "assistant": - continue - if message.get("tool_calls") and not message.get("content"): - continue - return index - return None - - def extract_attachment_ids(history_list) -> list[str]: attachment_ids = [] for history in history_list: @@ -487,7 +373,7 @@ class ChatRunState: run_id: str username: str session_id: str - llm_checkpoint_id: str + turn_id: str platform_history_id: str back_queue: asyncio.Queue subscribers: set[asyncio.Queue] = field(default_factory=set) @@ -672,60 +558,13 @@ async def delete_threads_by_ids(self, thread_ids: list[str], creator: str) -> No webchat_queue_mgr.remove_queues(thread_id) self.running_convs.pop(thread_id, None) - async def load_current_conversation_history(self, session) -> tuple[str, list]: - unified_msg_origin = build_webchat_unified_msg_origin(session) - conversation_id = await self.conv_mgr.get_curr_conversation_id( - unified_msg_origin - ) - if not conversation_id: - return "", [] - - conversation = await self.conv_mgr.get_conversation( - unified_msg_origin=unified_msg_origin, - conversation_id=conversation_id, - ) - if not conversation: - return "", [] - - try: - history = json.loads(conversation.history or "[]") - except json.JSONDecodeError: - return "", [] - return conversation_id, history if isinstance(history, list) else [] - - async def get_sorted_platform_history(self, session) -> list: - history_list = await self.platform_history_mgr.get( - platform_id=session.platform_id, - user_id=session.session_id, - page=1, - page_size=100000, - ) - history_list.sort(key=lambda item: (item.created_at, item.id)) - return history_list - - async def delete_platform_history_after( - self, session, message_id: int - ) -> list[int]: - history_list = await self.get_sorted_platform_history(session) - should_delete = False - deleted_ids: list[int] = [] - for item in history_list: - if should_delete: - if item.id is not None: - deleted_ids.append(item.id) - await self.platform_history_mgr.delete_by_id(item.id) - continue - if item.id == message_id: - should_delete = True - return deleted_ids - async def save_bot_message( self, webchat_conv_id: str, message_parts: list[dict], agent_stats: dict, refs: dict, - llm_checkpoint_id: str | None = None, + turn_id: str | None = None, platform_history_id: str = "webchat", ): return await self.platform_history_mgr.insert( @@ -738,7 +577,7 @@ async def save_bot_message( ), sender_id="bot", sender_name="bot", - llm_checkpoint_id=llm_checkpoint_id, + turn_id=turn_id, ) def get_active_chat_runs(self, username: str, session_id: str) -> list[dict]: @@ -759,7 +598,7 @@ def get_active_chat_runs(self, username: str, session_id: str) -> list[dict]: { "run_id": run.run_id, "session_id": run.session_id, - "llm_checkpoint_id": run.llm_checkpoint_id, + "turn_id": run.turn_id, "status": run.status, "revision": run.revision, "content": build_bot_history_content( @@ -817,7 +656,7 @@ def _subscribe_chat_run( snapshot = { "run_id": run.run_id, "session_id": run.session_id, - "llm_checkpoint_id": run.llm_checkpoint_id, + "turn_id": run.turn_id, "status": run.status, "revision": run.revision, "content": build_bot_history_content( @@ -848,7 +687,7 @@ async def stream(): "created_at": to_utc_isoformat( saved_user_record.created_at ), - "llm_checkpoint_id": run.llm_checkpoint_id, + "turn_id": run.turn_id, }, } yield f"data: {json.dumps(user_saved_info, ensure_ascii=False)}\n\n" @@ -934,7 +773,7 @@ async def flush_pending_bot_message(): message_parts_to_save, pending_agent_stats, extracted_refs, - run.llm_checkpoint_id, + run.turn_id, run.platform_history_id, ) pending_accumulator = BotMessageAccumulator() @@ -1035,7 +874,7 @@ async def flush_pending_bot_message(): "created_at": to_utc_isoformat( saved_record.created_at ), - "llm_checkpoint_id": run.llm_checkpoint_id, + "turn_id": run.turn_id, }, }, ) @@ -1062,7 +901,7 @@ async def flush_pending_bot_message(): "data": { "id": saved_record.id, "created_at": to_utc_isoformat(saved_record.created_at), - "llm_checkpoint_id": run.llm_checkpoint_id, + "turn_id": run.turn_id, }, }, ) @@ -1135,7 +974,7 @@ async def build_chat_stream( ) message_id = str(uuid.uuid4()) - llm_checkpoint_id = post_data.get("_llm_checkpoint_id") or str(uuid.uuid4()) + turn_id = post_data.get("_turn_id") or str(uuid.uuid4()) skip_user_history = bool(post_data.get("_skip_user_history")) saved_user_record = None @@ -1147,9 +986,30 @@ async def build_chat_stream( content={"type": "user", "message": message_parts_for_storage}, sender_id=username, sender_name=username, - llm_checkpoint_id=llm_checkpoint_id, + turn_id=turn_id, ) + if skip_user_history: + from sqlmodel import select + + from astrbot.core.db.po import PlatformMessageHistory + + async with self.db.get_db() as db_session: + saved_user_record = ( + await db_session.execute( + select(PlatformMessageHistory) + .where( + PlatformMessageHistory.turn_id == turn_id, + PlatformMessageHistory.user_id == webchat_conv_id, + PlatformMessageHistory.platform_id == platform_history_id, + PlatformMessageHistory.is_active.is_(True), + PlatformMessageHistory.content["type"].as_string() + == "user", + ) + .limit(1) + ) + ).scalar_one_or_none() + back_queue = webchat_queue_mgr.get_or_create_back_queue( message_id, webchat_conv_id, @@ -1158,7 +1018,7 @@ async def build_chat_stream( run_id=message_id, username=username, session_id=webchat_conv_id, - llm_checkpoint_id=llm_checkpoint_id, + turn_id=turn_id, platform_history_id=platform_history_id, back_queue=back_queue, ) @@ -1186,7 +1046,7 @@ async def build_chat_stream( "selected_model": selected_model, "flags": flags, "message_id": message_id, - "llm_checkpoint_id": llm_checkpoint_id, + "turn_id": turn_id, "thread_selected_text": thread_selected_text, "_api_key_allow_admin_role": post_data.get( "_api_key_allow_admin_role" @@ -1239,6 +1099,8 @@ async def delete_session_internal(self, session, username: str) -> None: tasks.append(run.task) if tasks: await asyncio.gather(*tasks, return_exceptions=True) + thread_ids = await self.db.delete_webchat_threads_by_parent_session(session_id) + await self.delete_threads_by_ids(thread_ids, username) await self.conv_mgr.delete_conversations_by_user_id(unified_msg_origin) history_list = await self.platform_history_mgr.get( @@ -1256,8 +1118,6 @@ async def delete_session_internal(self, session, username: str) -> None: user_id=session_id, offset_sec=99999999, ) - thread_ids = await self.db.delete_webchat_threads_by_parent_session(session_id) - await self.delete_threads_by_ids(thread_ids, username) try: await self.umop_config_router.delete_route(unified_msg_origin) @@ -1524,9 +1384,9 @@ async def create_thread(self, username: str, data: dict) -> dict: if parent_record.content.get("type") != "bot": raise ChatServiceError("Only bot messages can create threads") - checkpoint_id = parent_record.llm_checkpoint_id - if not checkpoint_id: - raise ChatServiceError("Parent message is not linked to LLM history") + source_event_id = parent_record.context_event_id + if not source_event_id: + raise ChatServiceError("Parent message is not linked to a context event") existing = await self.db.get_webchat_thread_by_parent_message_and_text( parent_session_id=session_id, @@ -1537,28 +1397,20 @@ async def create_thread(self, username: str, data: dict) -> dict: if existing: return serialize_thread(existing) - conversation_id, history = await self.load_current_conversation_history(session) - turn_range = find_turn_range(history, checkpoint_id) - if not conversation_id or not turn_range: - raise ChatServiceError("Linked checkpoint not found") - - _start, end = turn_range - base_history = history[: end + 1] thread = await self.db.create_webchat_thread( creator=username, parent_session_id=session_id, parent_message_id=parent_message_id, - base_checkpoint_id=checkpoint_id, + base_event_id=source_event_id, selected_text=selected_text, ) - await self.conv_mgr.new_conversation( - unified_msg_origin=build_thread_unified_msg_origin( - username, - thread.thread_id, - ), - platform_id="webchat", - content=base_history, - ) + thread_umo = build_thread_unified_msg_origin(username, thread.thread_id) + try: + cid = await self.conv_mgr.fork_conversation(thread_umo, source_event_id) + await self.conv_mgr.switch_conversation(thread_umo, cid) + except Exception: + await self.db.delete_webchat_thread(thread.thread_id) + raise return serialize_thread(thread) async def create_thread_from_dashboard_payload( @@ -1685,60 +1537,30 @@ async def update_message(self, username: str, data: dict) -> dict: if content.get("type") != "user": raise ChatServiceError("Only user messages can be edited") - platform_history = await self.get_sorted_platform_history(session) - latest_user_record = next( - ( - item - for item in reversed(platform_history) - if isinstance(item.content, dict) and item.content.get("type") == "user" - ), - None, + if self.running_convs.get(session_id, False): + raise ChatServiceError("Stop the active run before editing") + umo = build_webchat_unified_msg_origin(session) + conversation_id = await self.conv_mgr.get_curr_conversation_id(umo) + snapshot = ( + await self.db.conversation_store.read(conversation_id) + if conversation_id + else None ) - if not latest_user_record or latest_user_record.id != message_id: - raise ChatServiceError("Only the latest user message can be edited") - - checkpoint_id = record.llm_checkpoint_id - if not checkpoint_id: - raise ChatServiceError( - "This message is not linked to LLM history and cannot be edited" + if snapshot is None: + raise ChatServiceError("Conversation not found") + try: + replacement = await self.db.conversation_store.rewind_webchat( + conversation_id, + message_id, + content=content, + expected_head=snapshot.conversation.head_seq, + expected_leaf=snapshot.conversation.leaf_event_id, ) - - conversation_id, history = await self.load_current_conversation_history(session) - turn_range = find_turn_range(history, checkpoint_id) - if not conversation_id or not turn_range: - raise ChatServiceError("Linked checkpoint not found") - if not is_latest_checkpoint(history, checkpoint_id): - raise ChatServiceError("Only the latest turn can be edited") - - start, end = turn_range - target_index = find_turn_user_index(history, start, end) - if target_index is None: - raise ChatServiceError("Linked user message not found") - - new_checkpoint_id = str(uuid.uuid4()) - truncated_history = history[:start] - await self.platform_history_mgr.update( - message_id=message_id, - content=content, - llm_checkpoint_id=new_checkpoint_id, - ) - deleted_message_ids = await self.delete_platform_history_after( - session, message_id - ) - thread_ids = await self.db.delete_webchat_threads_by_parent_message_ids( - session_id, - deleted_message_ids, - ) - await self.delete_threads_by_ids(thread_ids, username) - await self.conv_mgr.update_conversation( - unified_msg_origin=build_webchat_unified_msg_origin(session), - conversation_id=conversation_id, - history=truncated_history, - ) + except ValueError as exc: + raise ChatServiceError(str(exc)) from exc await self.db.update_platform_session(session_id=session_id) - updated = await self.db.get_platform_message_history_by_id(message_id) return { - "message": serialize_history_entry(updated) if updated else None, + "message": serialize_history_entry(replacement), "needs_regenerate": True, "truncated_after_message": True, } @@ -1786,74 +1608,34 @@ async def prepare_regenerate_message_payload( if target_record.content.get("type") != "bot": raise ChatServiceError("Only bot messages can be regenerated") - checkpoint_id = target_record.llm_checkpoint_id - if not checkpoint_id: - raise ChatServiceError("Message is not linked to LLM history") - - conversation_id, history = await self.load_current_conversation_history(session) - turn_range = find_turn_range(history, checkpoint_id) - if not conversation_id or not turn_range: - raise ChatServiceError("Linked checkpoint not found") - if not is_latest_checkpoint(history, checkpoint_id): - raise ChatServiceError("Regenerating older turns requires branching") - - start, end = turn_range - user_index = find_turn_user_index(history, start, end) - if user_index is None: - raise ChatServiceError("Linked user message not found") - - platform_history = await self.get_sorted_platform_history(session) - source_user_record = next( - ( - item - for item in reversed(platform_history) - if item.llm_checkpoint_id == checkpoint_id - and isinstance(item.content, dict) - and item.content.get("type") == "user" - ), - None, + if self.running_convs.get(session_id, False): + raise ChatServiceError("Stop the active run before regenerating") + umo = build_webchat_unified_msg_origin(session) + conversation_id = await self.conv_mgr.get_curr_conversation_id(umo) + snapshot = ( + await self.db.conversation_store.read(conversation_id) + if conversation_id + else None ) - if not source_user_record: - raise ChatServiceError("Linked user display message not found") - - old_bot_record_ids = [ - item.id - for item in platform_history - if item.id is not None - and item.llm_checkpoint_id == checkpoint_id - and isinstance(item.content, dict) - and item.content.get("type") == "bot" - ] - if not old_bot_record_ids: - raise ChatServiceError("Linked bot display message not found") - - new_checkpoint_id = str(uuid.uuid4()) - new_history = history[:start] + history[end + 1 :] - await self.conv_mgr.update_conversation( - unified_msg_origin=build_webchat_unified_msg_origin(session), - conversation_id=conversation_id, - history=new_history, - ) - thread_ids = await self.db.delete_webchat_threads_by_parent_message_ids( - session_id, - old_bot_record_ids, - ) - await self.delete_threads_by_ids(thread_ids, username) - for old_bot_record_id in old_bot_record_ids: - await self.platform_history_mgr.delete_by_id(old_bot_record_id) - await self.platform_history_mgr.update( - message_id=source_user_record.id, - llm_checkpoint_id=new_checkpoint_id, - ) - + if snapshot is None: + raise ChatServiceError("Conversation not found") + try: + replacement = await self.db.conversation_store.rewind_webchat( + conversation_id, + message_id, + expected_head=snapshot.conversation.head_seq, + expected_leaf=snapshot.conversation.leaf_event_id, + ) + except ValueError as exc: + raise ChatServiceError(str(exc)) from exc return { "session_id": session_id, - "message": source_user_record.content.get("message", []), + "message": replacement.content.get("message", []), "flags": resolve_webchat_request_flags(data), "selected_provider": data.get("selected_provider"), "selected_model": data.get("selected_model"), "_skip_user_history": True, - "_llm_checkpoint_id": new_checkpoint_id, + "_turn_id": replacement.turn_id, } async def prepare_regenerate_message_payload_from_dashboard_payload( diff --git a/astrbot/dashboard/services/live_chat_service.py b/astrbot/dashboard/services/live_chat_service.py index 16b7eed0ad..99a53f102c 100644 --- a/astrbot/dashboard/services/live_chat_service.py +++ b/astrbot/dashboard/services/live_chat_service.py @@ -337,7 +337,7 @@ async def save_bot_message( message_parts: list[dict], agent_stats: dict, refs: dict, - llm_checkpoint_id: str | None = None, + turn_id: str | None = None, ): new_his = build_bot_history_content( message_parts, @@ -351,7 +351,7 @@ async def save_bot_message( content=new_his, sender_id="bot", sender_name="bot", - llm_checkpoint_id=llm_checkpoint_id, + turn_id=turn_id, ) async def send_chat_payload( @@ -556,7 +556,7 @@ async def handle_chat_message( await self.ensure_chat_subscription(session, session_id, send_json) back_queue = webchat_queue_mgr.get_or_create_back_queue(message_id, session_id) - llm_checkpoint_id = str(uuid.uuid4()) + turn_id = str(uuid.uuid4()) pending_bot_message_flusher = None try: @@ -575,7 +575,7 @@ async def handle_chat_message( "show_reasoning": show_reasoning, "flags": flags, "message_id": message_id, - "llm_checkpoint_id": llm_checkpoint_id, + "turn_id": turn_id, }, ), ) @@ -587,7 +587,7 @@ async def handle_chat_message( content={"type": "user", "message": message_parts_for_storage}, sender_id=session.username, sender_name=session.username, - llm_checkpoint_id=llm_checkpoint_id, + turn_id=turn_id, ) await self.send_chat_payload( session, @@ -597,7 +597,7 @@ async def handle_chat_message( "data": { "id": saved_user_record.id, "created_at": to_utc_isoformat(saved_user_record.created_at), - "llm_checkpoint_id": llm_checkpoint_id, + "turn_id": turn_id, }, **request_metadata, }, @@ -636,7 +636,7 @@ async def flush_pending_bot_message(): message_parts_to_save, agent_stats, extracted_refs, - llm_checkpoint_id, + turn_id, ) message_accumulator = BotMessageAccumulator() agent_stats = {} @@ -771,7 +771,7 @@ async def send_attachment_saved_event(part: dict | None) -> None: "created_at": to_utc_isoformat( saved_record.created_at ), - "llm_checkpoint_id": llm_checkpoint_id, + "turn_id": turn_id, }, **request_metadata, }, diff --git a/astrbot/dashboard/services/session_management_service.py b/astrbot/dashboard/services/session_management_service.py index 2307ae1adc..d9ba234f56 100644 --- a/astrbot/dashboard/services/session_management_service.py +++ b/astrbot/dashboard/services/session_management_service.py @@ -9,7 +9,7 @@ from astrbot.core import logger, sp from astrbot.core.core_lifecycle import AstrBotCoreLifecycle from astrbot.core.db import BaseDatabase -from astrbot.core.db.po import ConversationV2, Preference +from astrbot.core.db.po import ConversationV3, Preference from astrbot.core.provider.entities import ProviderType from astrbot.core.umo_alias import build_umo_alias_map, parse_umo, serialize_umo_alias @@ -57,7 +57,7 @@ def _is_private_umo(umo: str) -> bool: async def list_known_umos(self) -> list[str]: async with self.db_helper.get_db() as session: session: AsyncSession - result = await session.execute(select(ConversationV2.user_id).distinct()) + result = await session.execute(select(ConversationV3.umo).distinct()) umos = {str(row[0]) for row in result.fetchall() if row[0]} aliases = await self.db_helper.get_umo_aliases() diff --git a/dashboard/src/api/generated/openapi-v1/types.gen.ts b/dashboard/src/api/generated/openapi-v1/types.gen.ts index 4ac2fee95c..d16cb6e5ad 100644 --- a/dashboard/src/api/generated/openapi-v1/types.gen.ts +++ b/dashboard/src/api/generated/openapi-v1/types.gen.ts @@ -153,7 +153,7 @@ export type ChatRequest = { /** * Internal WebUI checkpoint override. */ - _llm_checkpoint_id?: string; + _turn_id?: string; /** * Internal WebUI platform history override. */ diff --git a/dashboard/src/assets/mdi-subset/materialdesignicons-subset.css b/dashboard/src/assets/mdi-subset/materialdesignicons-subset.css index f045ea72c8..dfd5b1db13 100644 --- a/dashboard/src/assets/mdi-subset/materialdesignicons-subset.css +++ b/dashboard/src/assets/mdi-subset/materialdesignicons-subset.css @@ -1,4 +1,4 @@ -/* Auto-generated MDI subset – 272 icons */ +/* Auto-generated MDI subset – 269 icons */ /* Do not edit manually. Run: pnpm run subset-icons */ @font-face { @@ -804,10 +804,6 @@ content: "\F1353"; } -.mdi-phone-in-talk::before { - content: "\F03F6"; -} - .mdi-pin::before { content: "\F0403"; } @@ -976,14 +972,6 @@ content: "\F04D2"; } -.mdi-stop::before { - content: "\F04DB"; -} - -.mdi-stop-circle::before { - content: "\F0666"; -} - .mdi-subdirectory-arrow-right::before { content: "\F060D"; } diff --git a/dashboard/src/assets/mdi-subset/materialdesignicons-webfont-subset.woff b/dashboard/src/assets/mdi-subset/materialdesignicons-webfont-subset.woff index bc7b4ddf1c..ec70455648 100644 Binary files a/dashboard/src/assets/mdi-subset/materialdesignicons-webfont-subset.woff and b/dashboard/src/assets/mdi-subset/materialdesignicons-webfont-subset.woff differ diff --git a/dashboard/src/assets/mdi-subset/materialdesignicons-webfont-subset.woff2 b/dashboard/src/assets/mdi-subset/materialdesignicons-webfont-subset.woff2 index 1ceed411c0..e362bd7fba 100644 Binary files a/dashboard/src/assets/mdi-subset/materialdesignicons-webfont-subset.woff2 and b/dashboard/src/assets/mdi-subset/materialdesignicons-webfont-subset.woff2 differ diff --git a/dashboard/src/components/chat/ChatMessageList.vue b/dashboard/src/components/chat/ChatMessageList.vue index ba092a8d75..11da5c9f98 100644 --- a/dashboard/src/components/chat/ChatMessageList.vue +++ b/dashboard/src/components/chat/ChatMessageList.vue @@ -607,37 +607,22 @@ function isEditingMessage(message: ChatRecord) { ); } -function canEditMessage(message: ChatRecord, messageIndex: number) { +function canEditMessage(message: ChatRecord, _messageIndex: number) { return ( props.enableEdit && isUserMessage(message) && - messageIndex === latestEditableUserIndex() && + Boolean(message.turn_id) && message.id != null && !String(message.id).startsWith("local-") ); } -function latestEditableUserIndex() { - for (let index = props.messages.length - 1; index >= 0; index -= 1) { - const message = props.messages[index]; - if ( - isUserMessage(message) && - message.id != null && - !String(message.id).startsWith("local-") - ) { - return index; - } - } - return -1; -} - function canRegenerateMessage(message: ChatRecord, messageIndex: number) { return ( props.enableRegenerate && !isUserMessage(message) && - messageIndex === props.messages.length - 1 && !isMessageStreaming(message, messageIndex) && - Boolean(message.llm_checkpoint_id) + Boolean(message.turn_id) ); } diff --git a/dashboard/src/components/chat/ThreadPanel.vue b/dashboard/src/components/chat/ThreadPanel.vue index 39d133fa5e..2becfed724 100644 --- a/dashboard/src/components/chat/ThreadPanel.vue +++ b/dashboard/src/components/chat/ThreadPanel.vue @@ -259,8 +259,8 @@ function processPayload( if (type === "user_message_saved") { userRecord.id = data?.id || userRecord.id; userRecord.created_at = data?.created_at || userRecord.created_at; - userRecord.llm_checkpoint_id = - data?.llm_checkpoint_id || userRecord.llm_checkpoint_id; + userRecord.turn_id = + data?.turn_id || userRecord.turn_id; return; } @@ -268,8 +268,8 @@ function processPayload( markMessageStarted(botRecord); botRecord.id = data?.id || botRecord.id; botRecord.created_at = data?.created_at || botRecord.created_at; - botRecord.llm_checkpoint_id = - data?.llm_checkpoint_id || botRecord.llm_checkpoint_id; + botRecord.turn_id = + data?.turn_id || botRecord.turn_id; if (data?.refs) { botRecord.content.refs = data.refs; } diff --git a/dashboard/src/composables/useMessages.ts b/dashboard/src/composables/useMessages.ts index 83068b1e59..9a65a79825 100644 --- a/dashboard/src/composables/useMessages.ts +++ b/dashboard/src/composables/useMessages.ts @@ -61,7 +61,8 @@ export interface ChatRecord { created_at?: string; sender_id?: string; sender_name?: string; - llm_checkpoint_id?: string | null; + turn_id?: string | null; + context_event_id?: string | null; threads?: ChatThread[]; } @@ -78,7 +79,7 @@ export interface ChatThread { thread_id: string; parent_session_id: string; parent_message_id: number; - base_checkpoint_id: string; + base_event_id: string; selected_text: string; created_at?: string; updated_at?: string; @@ -93,7 +94,7 @@ export interface ChatSessionProject { interface ActiveChatRun { run_id: string; session_id: string; - llm_checkpoint_id?: string | null; + turn_id?: string | null; status?: string; revision?: number; content?: ChatContent; @@ -128,7 +129,7 @@ interface SendMessageStreamOptions { userRecord?: ChatRecord; botRecord: ChatRecord; skipUserHistory?: boolean; - llmCheckpointId?: string | null; + turnId?: string | null; } interface ContinueEditedMessageOptions { @@ -434,18 +435,18 @@ export function useMessages(options: UseMessagesOptions) { const run = activeRuns[0]; if (!run?.run_id || isSessionRunning(sessionId)) return; - const checkpointId = run.llm_checkpoint_id || null; + const checkpointId = run.turn_id || null; const records = (messagesBySession[sessionId] || []).filter((record) => { return !( checkpointId && - record.llm_checkpoint_id === checkpointId && + record.turn_id === checkpointId && messageContent(record).type === "bot" ); }); const botRecord = normalizeHistoryRecord({ id: `active-run-${run.run_id}`, content: run.content || { type: "bot", message: [] }, - llm_checkpoint_id: checkpointId, + turn_id: checkpointId, created_at: new Date().toISOString(), }); botRecord.content.isLoading = botRecord.content.message.length === 0; @@ -525,7 +526,7 @@ export function useMessages(options: UseMessagesOptions) { botRecord, userRecord, skipUserHistory = false, - llmCheckpointId = null, + turnId = null, }: SendMessageStreamOptions) { if (transport === "websocket") { startWebSocketStream( @@ -552,7 +553,7 @@ export function useMessages(options: UseMessagesOptions) { selectedProvider, selectedModel, skipUserHistory, - llmCheckpointId, + turnId, ); } @@ -635,7 +636,7 @@ export function useMessages(options: UseMessagesOptions) { selectedProvider, selectedModel, true, - sourceRecord.llm_checkpoint_id || null, + sourceRecord.turn_id || null, ); } @@ -649,6 +650,12 @@ export function useMessages(options: UseMessagesOptions) { ) { if (!sessionId || botRecord.id == null) return; const targetMessageId = botRecord.id; + const records = messagesBySession[sessionId] || []; + const targetIndex = records.indexOf(botRecord); + const userRecord = records.slice(0, targetIndex).reverse().find( + (record) => record.content.type === "user" && record.turn_id === botRecord.turn_id, + ); + if (targetIndex >= 0) messagesBySession[sessionId] = records.slice(0, targetIndex + 1); botRecord.id = `local-regenerate-${Date.now()}`; botRecord.created_at = new Date().toISOString(); @@ -666,6 +673,7 @@ export function useMessages(options: UseMessagesOptions) { transport: "sse", abort, botRecord, + userRecord, botVisible: true, }; activeConnections[connection.messageId] = connection; @@ -780,7 +788,7 @@ export function useMessages(options: UseMessagesOptions) { selectedProvider: string, selectedModel: string, skipUserHistory = false, - llmCheckpointId: string | null = null, + turnId: string | null = null, ) { const abort = new AbortController(); const connection: ActiveConnection = { @@ -807,7 +815,7 @@ export function useMessages(options: UseMessagesOptions) { selected_provider: selectedProvider, selected_model: selectedModel, _skip_user_history: skipUserHistory, - _llm_checkpoint_id: llmCheckpointId || undefined, + _turn_id: turnId || undefined, }), signal: abort.signal, }) @@ -1205,13 +1213,13 @@ export function useMessages(options: UseMessagesOptions) { const snapshotRecord = normalizeHistoryRecord({ id: `active-run-${snapshot.run_id || "unknown"}`, content: snapshot.content || { type: "bot", message: [] }, - llm_checkpoint_id: snapshot.llm_checkpoint_id || null, + turn_id: snapshot.turn_id || null, }); snapshotRecord.content.isLoading = snapshot.status === "running" && snapshotRecord.content.message.length === 0; botRecord.content = snapshotRecord.content; - botRecord.llm_checkpoint_id = snapshotRecord.llm_checkpoint_id; + botRecord.turn_id = snapshotRecord.turn_id; void resolveRecordMedia([botRecord]); return; } @@ -1219,8 +1227,8 @@ export function useMessages(options: UseMessagesOptions) { if (userRecord) { userRecord.id = data?.id || userRecord.id; userRecord.created_at = data?.created_at || userRecord.created_at; - userRecord.llm_checkpoint_id = - data?.llm_checkpoint_id || userRecord.llm_checkpoint_id; + userRecord.turn_id = + data?.turn_id || userRecord.turn_id; } return; } @@ -1228,8 +1236,8 @@ export function useMessages(options: UseMessagesOptions) { markMessageStarted(botRecord); botRecord.id = data?.id || botRecord.id; botRecord.created_at = data?.created_at || botRecord.created_at; - botRecord.llm_checkpoint_id = - data?.llm_checkpoint_id || botRecord.llm_checkpoint_id; + botRecord.turn_id = + data?.turn_id || botRecord.turn_id; if (data?.refs) { messageContent(botRecord).refs = data.refs; } diff --git a/docs.md b/docs.md new file mode 100644 index 0000000000..1d69b1899f --- /dev/null +++ b/docs.md @@ -0,0 +1,399 @@ +# 新的 Agent 上下文存储架构 + +## 现状 + + + +AstrBot 默认使用 SQLite(`data_v4.db`),通过 SQLModel / SQLAlchemy 和 aiosqlite 访问,启用 WAL。 + + + +- 一个 UMO(`platform_name:message_type:session_id`)可以对应多个 Conversation,并选择当前对话。 + +- `ConversationV2.content` 保存整份模型消息数组。Manager 为兼容旧接口,又将其转换为 `Conversation.history` JSON 字符串。 + +- 内置 Agent 主要按“读取完整历史 → 在内存中修改、压缩 → 整体覆盖保存”工作。压缩后的内容会替代原历史。 + +- 平台可见消息另存于 `platform_message_history`。内部 checkpoint 用于关联平台消息与上下文轮次,不是执行快照;WebChat 侧边对话目前通过复制历史实现。 + +- 插件可以直接修改上下文;保存上下文时会过滤 `_no_save` 内容。执行状态主要在内存中,数据库写入缺少统一的并发版本校验。 + + + +## 目标与边界 + +以不可变事件保存消息与执行记录,按当前分支投影出模型需要的上下文。支持长期不切换对话的大事件量、消息编辑、侧边对话,以及旧插件兼容。 + +**事件 append\-only;ConversationV3 的当前指针和元数据可以更新。** 上下文恢复不重新调用模型、工具或插件,也不等于自动恢复外部工具执行。 + + + +## 表结构 + +### ConversationV3 + +|字段|类型|用途| +|---|---|---| +|`id`|INTEGER,主键|数据库内部标识| +|`conversation_id`|UUID,唯一|对外稳定标识| +|`platform_id`|TEXT|平台实例标识| +|`umo`|TEXT
|当前会话归属,替代含义不清的 `user_id`| +|`title`|TEXT,可空|对话标题| +|`persona_id`|TEXT,可空|当前人格配置| +|`head_seq`|BIGINT,默认 0|本对话已提交的最大事件序号| +|`leaf_event_id`|UUID,可空|当前选中分支的上下文末端| +|`replay_from_event_id`|UUID,可空|当前分支上可直接恢复的 `context.rebased` 节点| +|`created_at`|UTC 时间|创建时间| +|`updated_at`|UTC 时间|最近业务更新时间| + + + +不再保存 `content`;token 用量放在请求事件中。`umo` 暂时表达现有归属关系,每轮触发仍记录实际 UMO,为未来多 UMO 共用对话保留信息。 + + + +`replay_from_event_id` 是可重建的优化指针,不是另一份执行 checkpoint。使用全局事件 ID,是因为侧边对话可能继承另一 Conversation 的恢复起点,而 `seq` 只在各自对话内唯一。 + + + +### ConversationEvent + +所有用户、所有对话的事件共用一张表,通过 `conversation_ref` 隔离查询。 + +|字段|类型|用途| +|---|---|---| +|`conversation_ref`|INTEGER,外键|指向 `ConversationV3.id`| +|`seq`|BIGINT,正整数|对话内递增序号,不因清空上下文而重置| +|`event_id`|UUID,唯一|全局事件身份,写入重试复用同一个 ID| +|`parent_event_id`|UUID,可空,外键|上下文父节点,可跨 Conversation;空值表示无父节点| +|`type`|TEXT|事件类型| +|`version`|INTEGER,默认 1|对应类型的 payload 版本| +|`payload`|JSON|类型专属数据| +|`created_at`|UTC 时间|记录时间,排序以 seq 为准| + + + +约束与索引: + +- 主键 `(conversation_ref, seq)`,唯一索引 `event_id`。 + +- 外键保证引用存在;跨会话引用还需校验访问权限,不能级联删除被其他分支引用的事件。 + +- parent 只能引用已存在的上下文节点,或同批次中更早的上下文节点,避免形成环。 + +- 插件按类型查询需要时,增加 `(conversation_ref, type, seq)` 索引。不默认索引整个 payload。 + + + +## 分支、读取与并发 + +只有 `message.appended` 和 `context.rebased` 构成上下文 parent 链;会话、执行和插件私有事件不推进上下文 leaf,执行关联使用 payload 中的 ID。 + +```Plain Text +原分支:m1 → m2 → m3 → m4 +编辑后:m1 → m2 → m5 → m6 +``` + + + +编辑产生新节点,原分支保留。侧边 Conversation 的首个上下文节点可以指向来源 Conversation 的节点。`seq` 表示本对话的写入顺序,不表示当前分支顺序,不能直接按 seq 范围把所有消息拼起来。 + +写入 API 可以省略 parent,表示接到当前 leaf;**落库必须写入解析后的真实父节点,不能用 NULL 表示“上一个事件”**。显式 NULL 表示新根。 + +切换分支直接更新 `leaf_event_id` 和对应的恢复指针,不增加 `branch.selected` 事件。因此日志不记录用户每次选择分支的操作历史。 + + + +长对话的处理: + +1. 从当前 leaf 沿 parent 找到有效 rebase,批量或通过递归查询取得路径,再按正向顺序投影;避免每个父节点一次数据库往返。 + +2. rebase 保存完整有效上下文,恢复时不必读取其之前的祖先;旧历史仍可单独分页查询。 + +3. 按 token 预算压缩上下文;即使不需要摘要,也可以在重放事件数或字节量过大时追加内容不变的 rebase。 + +4. 按明确游标分页(如 `seq > after_seq`),不用不断扩大的 OFFSET,也不用记忆上次 LIMIT 来猜恢复范围。媒体使用资源引用。 + +5. 内存只缓存活跃分支的有效投影。不能因为有恢复指针,就假设永远没有过长的尾部路径。 + + + +序号分配、事件插入和指针更新在同一事务内完成。提交校验预期 `head_seq` 和 `leaf_event_id`;模型响应绑定开始时的分支,不能落到用户后来切换的分支。异步生成 rebase 也必须检查来源状态是否变化。冲突不能通过重新执行有副作用的插件或工具来解决。 + + + +## 事件类型 + +|类型|含义| +|---|---| +|`conversation.created`|创建会话,记录初始元数据和可选分支来源| +|`conversation.updated`|修改标题、人格等会话属性| +|`message.appended`|追加用户、assistant、tool 等模型消息| +|`context.rebased`|建立新的完整上下文投影起点| +|`turn.started`|一轮 Agent 处理开始| +|`turn.finished`|整轮处理结束| +|`request.started`|一次实际模型请求尝试开始| +|`request.finished`|请求结束,记录状态、usage、输出引用或错误| +|`tool.started`|工具准备执行,记录实际参数| +|`tool.finished`|工具执行结束,记录结果引用或错误| +|`plugin..`|插件私有记录,默认不参与模型上下文投影| + + + +一轮 turn 可以包含多次 request 和工具执行。turn、request、工具执行实例分别使用其 started 事件的 `event_id` 作为身份。重试创建新的请求或执行实例。 + + + +### 对话 + +实际 ID 使用 UUID,以下用短 ID 示意;后续例子省略重复的公共字段。 + +```JSON +{ + "conversation_ref": 42, + "seq": 1, + "event_id": "e1", + "parent_event_id": null, + "type": "conversation.created", + "version": 1, + "payload": { + "platform_id": "my_qq_bot", + "umo": "my_qq_bot:FriendMessage:user_123", + "persona_id": "default" + }, + "created_at": "2026-09-13T10:00:00Z" +} +``` + + + +创建侧边对话时,payload 可增加 `forked_from_event_id`。更新属性时,字段缺席表示不修改,显式 null 表示清空: + +```JSON +{ + "type": "conversation.updated", + "payload": {"changes": {"title": "旅行计划", "persona_id": null}} +} +``` + + + +### 消息与 rebase + + + +消息内容沿用 AstrBot 消息模型,保留工具调用、多模态内容以及供应商需要的额外字段。下面以常见消息格式示意: + +```JSON +[ + { + "event_id": "m1", + "parent_event_id": null, + "type": "message.appended", + "payload": { + "turn_id": "t1", + "message": {"role": "user", "content": "查一下北京天气"} + } + }, + { + "event_id": "m2", + "parent_event_id": "m1", + "type": "message.appended", + "payload": { + "turn_id": "t1", + "request_id": "r1", + "message": { + "role": "assistant", + "content": null, + "tool_calls": [{ + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{\"city\":\"北京\"}"} + }] + } + } + }, + { + "event_id": "m3", + "parent_event_id": "m2", + "type": "message.appended", + "payload": { + "turn_id": "t1", + "message": {"role": "tool", "tool_call_id": "call_1", "content": "北京晴,25°C"} + } + } +] +``` + + + +`context.rebased` 统一使用 `reason` 区分 `compaction`、`reset`、`legacy_replace`、`migration` 和 `snapshot`。messages 是完整有效列表;如有其他继续投影所需的状态,也必须随事件保存。仅保存旧消息 ID 而需要扫描古老历史补齐内容,不能算完整恢复起点。 + + + +```JSON +{ + "event_id": "b1", + "parent_event_id": "m100", + "type": "context.rebased", + "payload": { + "reason": "compaction", + "messages": [ + {"id": "mxxx", "message": {"role": "system", "content": "xxxxxx"}}, + {"id": "summary_b1", "message": {"role": "user", "content": "此前摘要:用户计划周末去北京。"}}, + {"id": "m99", "message": {"role": "user", "content": "如果下雨呢?"}}, + {"id": "m100", "message": {"role": "assistant", "content": "可以参观室内博物馆。"}} + ] + } +} +``` + + + +保留消息沿用消息身份,新摘要分配新身份;消息 ID 不一定是独立事件外键。摘要角色与顺序由压缩策略决定。reset 使用空列表;snapshot 保持消息内容和顺序不变。优先在无待处理工具调用的稳定边界生成 rebase。 + + + +### turn、request 与 tool + + + +下面是各类型的独立示例,不是完整执行顺序: + +```JSON +[ + { + "event_id": "t1", + "type": "turn.started", + "payload": { + "trigger": {"kind": "im_wake", "umo": "my_qq_bot:FriendMessage:user_123", "message_id": "platform_msg_1"}, + "base_leaf_event_id": null + } + }, + { + "type": "turn.finished", + "payload": {"turn_id": "t1", "status": "completed"} + }, + { + "event_id": "r1", + "type": "request.started", + "payload": { + "turn_id": "t1", + "context_leaf_event_id": "m1", + "provider_id": "configured_provider_1", + "model": "configured_model", + "parameters": {"temperature": 0.7} + } + }, + { + "type": "request.finished", + "payload": { + "request_id": "r1", + "status": "completed", + "output_event_ids": ["m2"], + "finish_reason": "tool_calls", + "usage": {"input_tokens": 1200, "output_tokens": 80, "cached_input_tokens": 600} + } + }, + { + "event_id": "x1", + "type": "tool.started", + "payload": { + "turn_id": "t1", + "tool_call_id": "call_1", + "tool_name": "get_weather", + "arguments": {"city": "北京", "units": "celsius"} + } + }, + { + "type": "tool.finished", + "payload": {"execution_id": "x1", "status": "completed", "result_event_id": "m3"} + } +] +``` + + + +执行约定: + +- turn\.trigger 按来源提供字段,可为 `im_wake`、`plugin`、`agent`、`background_task_finished` 等,不填无意义的空字段。 + +- finished 用 `status` 区分 `completed`、`failed`、`cancelled`;失败可附 `error: {code, message}`。只有 started 表示结果未知,不能据此自动重试外部操作。 + +- `request.finished.usage` 是用量权威明细;turn 可保存汇总。 + +- `tool.started` 在执行前提交,记录插件 hook 处理后的实际参数。模型提出的 tool\_calls 不等于工具实际执行。 + +- finished 和对应消息可一起提交;执行结果与模型可见结果相同时使用引用,发生改写时分别保留必要内容,避免错误引用或重复大正文。 + +- 不默认逐条存储 token delta、工具进度和 hook 通知。崩溃恢复外部副作用还需要幂等或结果查询机制。 + +- `_no_save` 定义为:**内容照旧落盘,但是通过字段控制不进入投影。** + + ```JSON + { + "type": "message.appended", + "payload": { + **"include_in_context": false,** + "message": { + "role": "user", + "content": "仅供本次请求使用的临时提示" + } + } + } + ``` + +## 插件接口与兼容 + + + +旧插件继续使用原有的可变 `ProviderRequest.contexts`、运行时消息对象和 ConversationManager 接口。核心事件保持不可变,插件拿到与持久化状态隔离的工作副本;同一阶段的插件按顺序看到前面插件的修改。 + + + +兼容层在原有保存边界比较阶段输入与最终输出,不对每个插件复制整份历史: + + + +- 无变化:不产生上下文事件。 + +- 原消息前缀不变、只追加:生成 `message.appended`。 + +- 删除、重排、嵌套修改或整体替换:生成 `context.rebased(reason=legacy_replace)`。 + + + +插件在已绑定会话的 hook 中,通过 ConversationManager 获取统一接口: + +```Python +events = ctx.conversation_manager.get_conversation_events() +await events.append_message({"role": "user", "content": "补充上下文"}) +await events.append( + name="retrieval_completed", + payload={"document_ids": ["doc_12", "doc_35"]}, +) +previous = await events.latest(name="retrieval_completed") +page = await events.list(name="retrieval_completed", after_seq=last_seq, limit=100) +``` + + + +框架绑定插件身份与会话、分配 ID 和 seq、校验 JSON 并落库,生成例如: + +```JSON +{ + "event_id": "p1", + "type": "plugin.astrbot_plugin_memory.retrieval_completed", + "payload": {"turn_id": "t1", "document_ids": ["doc_12", "doc_35"], "scores": [0.91, 0.86]} +} +``` + + + +插件日志默认按当前 Conversation 和插件身份查询,不自动继承来源分支的插件状态;分支相关记录显式关联上下文事件 ID。插件私有记录不驱动核心消息投影,卸载插件也不影响核心上下文恢复。 + + + +想永久改变上下文,应调用核心追加消息或替换上下文接口;只给当前请求提供检索内容,则使用明确的临时请求接口。恢复时读取已记录的结果,不重新执行插件业务逻辑。 + + diff --git a/docs/public/openapi.json b/docs/public/openapi.json index b666e92afa..54ad6a7860 100644 --- a/docs/public/openapi.json +++ b/docs/public/openapi.json @@ -7201,7 +7201,7 @@ "type": "boolean", "description": "Internal WebUI flag for edit/regenerate flows." }, - "_llm_checkpoint_id": { + "_turn_id": { "type": "string", "description": "Internal WebUI checkpoint override." }, diff --git a/openspec/openapi-v1.yaml b/openspec/openapi-v1.yaml index 51a478f04b..f94ee276e3 100644 --- a/openspec/openapi-v1.yaml +++ b/openspec/openapi-v1.yaml @@ -5602,7 +5602,7 @@ components: _skip_user_history: type: boolean description: Internal WebUI flag for edit/regenerate flows. - _llm_checkpoint_id: + _turn_id: type: string description: Internal WebUI checkpoint override. _platform_history_id: diff --git a/tests/test_backup.py b/tests/test_backup.py index 52f60a48c1..adffa85896 100644 --- a/tests/test_backup.py +++ b/tests/test_backup.py @@ -1014,7 +1014,8 @@ def test_main_db_models_contain_expected_tables(self): """测试主数据库模型映射包含预期的表""" expected_tables = [ "platform_stats", - "conversations", + "conversations_v3", + "conversation_events", "personas", "preferences", "chatui_projects", diff --git a/tests/test_conversation_checkpoint.py b/tests/test_conversation_checkpoint.py index e063886767..678a3e3f68 100644 --- a/tests/test_conversation_checkpoint.py +++ b/tests/test_conversation_checkpoint.py @@ -22,7 +22,6 @@ ) from astrbot.core.provider.entities import LLMResponse, ProviderRequest, ToolCallsResult from astrbot.core.provider.provider import Provider -from astrbot.dashboard.services.chat_service import find_turn_range def test_checkpoint_message_segment_round_trip(): @@ -161,28 +160,16 @@ def test_provider_ensure_message_to_dicts_skips_checkpoints(): ] -def test_chat_service_find_turn_range(): - history = [ - {"role": "user", "content": "a"}, - {"role": "assistant", "content": "b"}, - {"role": "_checkpoint", "content": {"id": "cp-1"}}, - {"role": "user", "content": "c"}, - {"role": "assistant", "content": "d"}, - {"role": "_checkpoint", "content": {"id": "cp-2"}}, - ] - - assert find_turn_range(history, "cp-2") == (3, 5) - assert find_turn_range(history, "missing") is None @pytest.mark.asyncio -async def test_failed_llm_response_persists_checkpoint_for_retry(): +async def test_failed_llm_response_preserves_input_for_linked_turn_retry(): conversation_manager = AsyncMock() stage = InternalAgentSubStage() stage.conv_manager = conversation_manager event = SimpleNamespace( unified_msg_origin="webchat:FriendMessage:test", - get_extra=lambda key: {"llm_checkpoint_id": "cp-1"}.get(key), + get_extra=lambda key: {"turn_id": "cp-1"}.get(key), ) request = ProviderRequest( conversation=Conversation( @@ -205,7 +192,6 @@ async def test_failed_llm_response_persists_checkpoint_for_retry(): "conversation-1", history=[ {"role": "user", "content": "hello"}, - {"role": "_checkpoint", "content": {"id": "cp-1"}}, ], token_usage=None, ) @@ -332,13 +318,13 @@ async def test_terminal_tool_result_persists_history_without_checkpoint(): @pytest.mark.asyncio -async def test_terminal_tool_result_with_checkpoint_uses_none_token_usage(): +async def test_terminal_tool_result_with_turn_uses_none_token_usage(): conversation_manager = AsyncMock() stage = InternalAgentSubStage() stage.conv_manager = conversation_manager event = SimpleNamespace( unified_msg_origin="qq:GroupMessage:test", - get_extra=lambda key: {"llm_checkpoint_id": "cp-1"}.get(key), + get_extra=lambda key: {"turn_id": "cp-1"}.get(key), ) tool_call = ToolCall( id="call-1", @@ -398,7 +384,6 @@ async def test_terminal_tool_result_with_checkpoint_uses_none_token_usage(): "content": "The tool has no return value.", "tool_call_id": "call-1", }, - {"role": "_checkpoint", "content": {"id": "cp-1"}}, ], token_usage=None, ) diff --git a/tests/test_conversation_list.py b/tests/test_conversation_list.py index e12cbc3539..449860d20c 100644 --- a/tests/test_conversation_list.py +++ b/tests/test_conversation_list.py @@ -4,13 +4,32 @@ import pytest from sqlalchemy import event, text -from sqlalchemy import inspect as sqlalchemy_inspect from astrbot.core.conversation_mgr import ConversationManager -from astrbot.core.db.po import ConversationV2, PlatformSession +from astrbot.core.db.po import ConversationV2, ConversationV3, PlatformSession from astrbot.core.db.sqlite import SQLiteDatabase +async def seed_conversations(db, rows): + """Seed current storage through its public API, plus display session fixtures. + + Args: + db: Test database. + rows: Legacy-shaped conversation or platform session fixtures. + """ + for row in rows: + if isinstance(row, ConversationV2): + await db.create_conversation( + user_id=row.user_id, platform_id=row.platform_id, + content=row.content, title=row.title, persona_id=row.persona_id, + cid=row.conversation_id, created_at=row.created_at, updated_at=row.updated_at, + ) + else: + async with db.get_db() as session: + session.add(row) + await session.commit() + + @pytest.mark.asyncio async def test_filtered_conversations_summary_skips_content_and_applies_filters( tmp_path: Path, @@ -61,9 +80,7 @@ async def test_filtered_conversations_summary_skips_content_and_applies_filters( updated_at=datetime(2023, 1, 1, tzinfo=timezone.utc), ), ] - async with db.get_db() as session: - async with session.begin(): - session.add_all(conversations) + await seed_conversations(db, conversations) summary, total = await db.get_filtered_conversations( page=1, @@ -80,7 +97,7 @@ async def test_filtered_conversations_summary_skips_content_and_applies_filters( "group", "other", ] - assert all("content" in sqlalchemy_inspect(item).unloaded for item in summary) + assert all(item.content is None for item in summary) manager_summary, manager_total = await ConversationManager( db, @@ -120,7 +137,7 @@ async def test_filtered_conversations_summary_skips_content_and_applies_filters( full, full_total = await db.get_filtered_conversations(page_size=10) assert full_total == 5 - assert all("content" not in sqlalchemy_inspect(item).unloaded for item in full) + assert all(item.content is not None for item in full) umo_matches, _ = await db.get_filtered_conversations( umo_query="FriendMessage:2", @@ -177,17 +194,13 @@ def conversation(cid: str, user_id: str, day: int) -> ConversationV2: updated_at=timestamp, ) - async with db.get_db() as session: - async with session.begin(): - session.add_all( - [ + await seed_conversations(db, [ conversation("a-old", "qq:FriendMessage:a", 1), conversation("a-new", "qq:FriendMessage:a", 2), conversation("b-old", "qq:FriendMessage:b", 3), conversation("b-new", "qq:FriendMessage:b", 4), conversation("c-only", "qq:FriendMessage:c", 5), - ] - ) + ]) first_page, total_sessions = await db.get_filtered_conversations( page=1, @@ -213,7 +226,7 @@ def conversation(cid: str, user_id: str, day: int) -> ConversationV2: "b-old", ] assert [item.conversation_id for item in second_page] == ["a-new", "a-old"] - assert all("content" in sqlalchemy_inspect(item).unloaded for item in first_page) + assert all(item.content is None for item in first_page) @pytest.mark.asyncio @@ -226,26 +239,26 @@ async def test_conversation_indexes_are_idempotent_and_support_ordered_list( async with db.get_db() as session: index_rows = ( - await session.execute(text("PRAGMA index_list(conversations)")) + await session.execute(text("PRAGMA index_list(conversations_v3)")) ).all() index_names = {row[1] for row in index_rows} plan = ( await session.execute( text( "EXPLAIN QUERY PLAN " - "SELECT conversation_id FROM conversations " - "ORDER BY created_at DESC, inner_conversation_id DESC LIMIT 20" + "SELECT conversation_id FROM conversations_v3 " + "ORDER BY created_at DESC, id DESC LIMIT 20" ) ) ).all() expected_indexes = { - "ix_conversations_created_at_inner_id", - "ix_conversations_platform_created_at_inner_id", + "ix_conversations_v3_created_id", + "ix_conversations_v3_platform_created_id", } assert expected_indexes.issubset(index_names) assert expected_indexes.issubset( - {index.name for index in ConversationV2.__table__.indexes} + {index.name for index in ConversationV3.__table__.indexes} ) assert "ix_conversations_platform_user_id" not in index_names assert not any("TEMP B-TREE" in str(row) for row in plan) @@ -258,10 +271,7 @@ async def test_multi_platform_summary_uses_global_order_index( db = SQLiteDatabase(str(tmp_path / "multi-platform.db")) await db.initialize() - async with db.get_db() as session: - async with session.begin(): - session.add_all( - [ + await seed_conversations(db, [ ConversationV2( conversation_id=f"conversation-{index}", platform_id="qq" if index % 2 else "telegram", @@ -270,8 +280,7 @@ async def test_multi_platform_summary_uses_global_order_index( created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), ) for index in range(20) - ], - ) + ]) statements = [] @@ -301,12 +310,12 @@ def capture_statement(_conn, _cursor, statement, _parameters, _context, _many): "conversation-16", "conversation-15", ] - assert all("content" in sqlalchemy_inspect(item).unloaded for item in conversations) + assert all(item.content is None for item in conversations) - ordered_queries = [statement for statement in statements if "ORDER BY" in statement] + ordered_queries = [statement for statement in statements if "ORDER BY" in statement and "FROM conversations_v3" in statement] assert len(ordered_queries) == 1 assert ( - "FROM conversations INDEXED BY ix_conversations_created_at_inner_id" + "FROM conversations_v3 INDEXED BY ix_conversations_v3_created_id" in ordered_queries[0] ) assert "content" not in ordered_queries[0].split("FROM", 1)[0] @@ -319,10 +328,7 @@ async def test_webchat_session_title_matches_search_and_keyword(tmp_path: Path): await db.initialize() matched_session_id = "session-with-title" - async with db.get_db() as session: - async with session.begin(): - session.add_all( - [ + await seed_conversations(db, [ ConversationV2( conversation_id="webchat-titled", platform_id="webchat", @@ -354,8 +360,7 @@ async def test_webchat_session_title_matches_search_and_keyword(tmp_path: Path): creator="astrbot", display_name="成都旅行三日游计划", ), - ] - ) + ]) conversations, total = await db.get_filtered_conversations( page=1, diff --git a/tests/test_conversation_v3.py b/tests/test_conversation_v3.py new file mode 100644 index 0000000000..d80ad651c0 --- /dev/null +++ b/tests/test_conversation_v3.py @@ -0,0 +1,562 @@ +"""Behavioral checks for event storage, migration, and runner compatibility.""" + +import asyncio +import json +from pathlib import Path + +import pytest +import pytest_asyncio +from sqlalchemy import text +from sqlmodel import select + +from astrbot.core.agent.conversation_events import ConversationEventWriter +from astrbot.core.db.conversation import ConversationConflictError +from astrbot.core.db.po import ConversationEvent, ConversationV2, ConversationV3, PlatformMessageHistory +from astrbot.core.db.sqlite import SQLiteDatabase + + +@pytest_asyncio.fixture +async def database(tmp_path: Path): + db = SQLiteDatabase(str(tmp_path / "events.db")) + await db.initialize() + yield db + await db.engine.dispose() + + +@pytest.mark.asyncio +async def test_legacy_api_appends_and_rebases_without_losing_old_events(database): + conv = await database.create_conversation("umo", "platform", [{"role": "user", "content": "old"}]) + await database.update_conversation(conv.conversation_id, content=[ + {"role": "user", "content": "old"}, {"role": "assistant", "content": "reply"}]) + before = await database.conversation_store.events(conv.conversation_id) + assert before[-1].type == "message.appended" + await database.update_conversation(conv.conversation_id, content=[{"role": "user", "content": "edited"}]) + after = await database.conversation_store.events(conv.conversation_id) + assert [e.model_dump() for e in after[:len(before)]] == [e.model_dump() for e in before] + assert after[-1].type == "context.rebased" + restored = await database.get_conversation_by_id(conv.conversation_id) + assert restored.content == [{"role": "user", "content": "edited"}] + restored.content[0]["content"] = "local mutation" + assert (await database.get_conversation_by_id(conv.conversation_id)).content[0]["content"] == "edited" + + +@pytest.mark.asyncio +async def test_branches_null_parent_and_excluded_messages(database): + conv = await database.create_conversation("umo", "p") + store = database.conversation_store + first = (await store.append(conv.conversation_id, [{"event_id": "root", "type": "message.appended", "payload": {"message": {"role": "user", "content": "root"}}}]))[0] + await store.append(conv.conversation_id, [{"type": "message.appended", "payload": {"include_in_context": False, "message": {"role": "user", "content": "temporary"}}}]) + await store.append(conv.conversation_id, [{"type": "message.appended", "payload": {"message": {"role": "assistant", "content": "reply"}}}]) + assert [m["content"] for m in (await store.read(conv.conversation_id)).messages] == ["root", "reply"] + branch = await store.create(umo="umo", platform_id="p", parent_event_id=first.event_id) + await store.append(branch.conversation_id, [{"type": "message.appended", "payload": {"message": {"role": "assistant", "content": "branch"}}}]) + assert [m["content"] for m in (await store.read(branch.conversation_id)).messages] == ["root", "branch"] + with pytest.raises(ValueError, match="referenced"): + await store.delete(cid=conv.conversation_id) + await store.append(branch.conversation_id, [{"type": "message.appended", "parent_event_id": None, "payload": {"message": {"role": "user", "content": "new root"}}}]) + assert (await store.read(branch.conversation_id)).messages == [{"role": "user", "content": "new root"}] + with pytest.raises(ValueError, match="accessible"): + await store.create(umo="other", platform_id="p", parent_event_id=first.event_id) + + +@pytest.mark.asyncio +async def test_concurrent_writers_and_idempotent_delivery(database): + conv = await database.create_conversation("umo", "p") + store = database.conversation_store + snapshot = await store.read(conv.conversation_id) + writers = [ConversationEventWriter(store, snapshot) for _ in range(2)] + results = await asyncio.gather(*(w.append("plugin.test.result", {"value": i}, event_id=f"p{i}") for i, w in enumerate(writers)), return_exceptions=True) + assert sum(isinstance(r, ConversationConflictError) for r in results) == 1 + succeeded = next(i for i, r in enumerate(results) if not isinstance(r, Exception)) + repeated = await writers[succeeded].append("plugin.test.result", {"value": succeeded}, event_id=f"p{succeeded}") + assert repeated.event_id == f"p{succeeded}" + with pytest.raises(ConversationConflictError): + await writers[succeeded].append("plugin.test.result", {"value": "different"}, event_id=f"p{succeeded}") + assert len(await store.events(conv.conversation_id)) == 2 + + +@pytest.mark.asyncio +async def test_staged_plugin_changes_no_save_and_private_journal(database): + conv = await database.create_conversation("umo", "p") + writer = ConversationEventWriter(database.conversation_store, await database.conversation_store.read(conv.conversation_id)) + await writer.start_turn({"kind": "plugin"}) + writer.stage_history([{"role": "user", "content": "secret", "_no_save": True}, {"role": "user", "content": "old"}]) + writer.stage_history([{"role": "user", "content": "summary"}], reason="compaction") + assert (await database.get_conversation_by_id(conv.conversation_id)).content == [] + plugin = writer.plugin("memory") + await plugin.append("retrieval_finished", {"document_ids": ["d1"]}) + assert (await plugin.latest("retrieval_finished")).payload["document_ids"] == ["d1"] + await writer.save_history([{"role": "user", "content": "summary"}]) + await writer.finish_turn("completed") + events = await database.conversation_store.events(conv.conversation_id) + serialized = json.dumps([e.payload for e in events]) + assert "secret" in serialized and "old" in serialized + assert next(e for e in events if e.payload.get("message", {}).get("content") == "secret").payload["include_in_context"] is False + assert (await database.get_conversation_by_id(conv.conversation_id)).content == [{"role": "user", "content": "summary"}] + + +@pytest.mark.asyncio +async def test_replay_ignores_large_execution_log_and_pages_by_cursor(database): + conv = await database.create_conversation("umo", "p", [{"role": "user", "content": "kept"}]) + async with database.get_db() as session: + await session.execute(ConversationEvent.__table__.insert(), [ + {"conversation_ref": conv.inner_conversation_id, "seq": i + 3, "event_id": f"log{i}", "type": "plugin.test.progress", "payload": {"i": i}} + for i in range(100_000) + ]) + metadata = await session.get(ConversationV3, conv.inner_conversation_id) + metadata.head_seq = 100_002 + await session.commit() + snapshot = await database.conversation_store.read(conv.conversation_id) + assert snapshot.replay_count == 1 + assert snapshot.messages == [{"role": "user", "content": "kept"}] + page = await database.conversation_store.events(conv.conversation_id, after_seq=99_998, limit=3) + assert [e.seq for e in page] == [99_999, 100_000, 100_001] + + +@pytest.mark.asyncio +async def test_plugin_manager_entrypoint_isolates_concurrent_hooks(database): + from astrbot.core.agent.conversation_events import ( + active_conversation_writer, + active_plugin_id, + ) + from astrbot.core.agent.run_context import ContextWrapper + from astrbot.core.conversation_mgr import ConversationManager + + manager = ConversationManager(database) + store = database.conversation_store + writers = [] + for plugin_id in ("memory", "search"): + conv = await database.create_conversation(plugin_id, "p") + writer = await manager.event_writer(plugin_id, conv.conversation_id) + writer.runtime_context = ContextWrapper(context=None, messages=[]) + writers.append(writer) + + async def run_hook(writer, plugin_id): + writer_token = active_conversation_writer.set(writer) + plugin_token = active_plugin_id.set(plugin_id) + try: + await asyncio.sleep(0) + events = manager.get_conversation_events() + saved = await events.append("retrieval_completed", {"source": plugin_id}) + await events.append_message( + {"role": "user", "content": plugin_id}, event_id=plugin_id + ) + await events.append_message( + {"role": "user", "content": "temporary"}, include_in_context=False + ) + assert (await events.latest("retrieval_completed")).event_id == saved.event_id + assert len(await events.list("retrieval_completed")) == 1 + assert await events.list("retrieval_completed", after_seq=saved.seq) == [] + finally: + active_plugin_id.reset(plugin_token) + active_conversation_writer.reset(writer_token) + + await asyncio.gather( + *(run_hook(writer, plugin_id) for writer, plugin_id in zip(writers, ("memory", "search"))) + ) + for writer, plugin_id in zip(writers, ("memory", "search")): + snapshot = await store.read(writer.cid) + assert snapshot.messages == [{"role": "user", "content": plugin_id}] + private = await store.events(writer.cid, event_type=f"plugin.{plugin_id}.retrieval_completed") + assert len(private) == 1 + assert private[0].payload == {"source": plugin_id} + assert writer.runtime_context.messages[0].content == plugin_id + assert writer.runtime_context.messages[1]._no_save + + with pytest.raises(RuntimeError, match="active conversation hook"): + manager.get_conversation_events() + + writer_token = active_conversation_writer.set(writers[0]) + try: + with pytest.raises(RuntimeError, match="active conversation hook"): + manager.get_conversation_events() + finally: + active_conversation_writer.reset(writer_token) + + +@pytest.mark.asyncio +async def test_migration_drops_old_table_and_preserves_checkpoint_links(tmp_path): + db = SQLiteDatabase(str(tmp_path / "migration.db")) + async with db.engine.begin() as connection: + await connection.run_sync(ConversationV2.__table__.create) + await connection.execute(text("CREATE TABLE platform_message_history (id INTEGER PRIMARY KEY, platform_id TEXT NOT NULL, user_id TEXT NOT NULL, sender_id TEXT, sender_name TEXT, content JSON NOT NULL, llm_checkpoint_id TEXT, created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL)")) + history = [{"role": "user", "content": "hello"}, {"role": "assistant", "content": "world"}, {"role": "_checkpoint", "content": {"id": "cp1"}}] + async with db.AsyncSessionLocal() as session: + session.add(ConversationV2(conversation_id="old", platform_id="webchat", user_id="webchat:FriendMessage:webchat!alice!session1", content=history, token_usage=123)) + await session.execute(text("INSERT INTO platform_message_history VALUES (1, 'webchat', 'session1', NULL, NULL, :content, 'cp1', '2026-01-01', '2026-01-01')"), {"content": json.dumps({"type": "bot", "message": []})}) + await session.commit() + await db.initialize() + await db.initialize() + restored = await db.get_conversation_by_id("old") + assert restored.content == history[:2] + assert restored.token_usage == 123 + async with db.get_db() as session: + assert not (await session.execute(text("SELECT 1 FROM sqlite_master WHERE name='conversations'"))).first() + record = await session.get(PlatformMessageHistory, 1) + event = (await session.execute(select(ConversationEvent).where(ConversationEvent.event_id == record.context_event_id))).scalar_one() + assert event.payload["message"]["content"] == "world" + assert record.turn_id == event.payload["turn_id"] + await db.engine.dispose() + + +@pytest.mark.asyncio +async def test_webchat_edit_preserves_original_branch_and_side_thread(database): + store = database.conversation_store + umo = "webchat:FriendMessage:webchat!alice!session1" + conv = await store.create(umo=umo, platform_id="webchat") + writer = ConversationEventWriter(store, await store.read(conv.conversation_id)) + await writer.start_turn({"kind": "im_wake"}, event_id="t1") + user = await database.insert_platform_message_history("webchat", "session1", {"type": "user", "message": []}, turn_id="t1") + await writer.save_history([{"role": "user", "content": "old"}, {"role": "assistant", "content": "answer"}]) + bot = await database.insert_platform_message_history("webchat", "session1", {"type": "bot", "message": []}, turn_id="t1") + await writer.finish_turn("completed") + side = await store.create(umo="webchat:FriendMessage:webchat!alice!side", platform_id="webchat", parent_event_id=bot.context_event_id) + snapshot = await store.read(conv.conversation_id) + replacement = await store.rewind_webchat(conv.conversation_id, user.id, content={"type": "user", "message": [{"type": "plain", "text": "edited"}]}, expected_head=snapshot.conversation.head_seq, expected_leaf=snapshot.conversation.leaf_event_id) + assert replacement.id != user.id + assert (await store.read(conv.conversation_id)).messages == [] + assert [m["content"] for m in (await store.read(side.conversation_id)).messages] == ["old", "answer"] + visible = await database.get_platform_message_history("webchat", "session1") + assert [r.id for r in visible] == [replacement.id] + assert (await database.get_platform_message_history_by_id(user.id)).is_active is False + + +@pytest.mark.asyncio +async def test_failed_migration_keeps_legacy_data_and_can_retry(tmp_path, monkeypatch): + db = SQLiteDatabase(str(tmp_path / "rollback.db")) + async with db.engine.begin() as connection: + await connection.run_sync(ConversationV2.__table__.create) + original = [{"role": "user", "content": "must survive"}] + async with db.AsyncSessionLocal() as session: + session.add(ConversationV2(conversation_id="old", platform_id="p", user_id="umo", content=original)) + await session.commit() + project = db.conversation_store.project + + async def fail_verification(*args, **kwargs): + raise ValueError("injected verification failure") + + monkeypatch.setattr(db.conversation_store, "project", fail_verification) + with pytest.raises(ValueError, match="injected"): + await db.initialize() + assert not db.inited + async with db.AsyncSessionLocal() as session: + old = (await session.execute(select(ConversationV2))).scalar_one() + assert old.content == original + assert not (await session.execute(select(ConversationV3))).first() + assert not (await session.execute(select(ConversationEvent))).first() + monkeypatch.setattr(db.conversation_store, "project", project) + await db.initialize() + assert (await db.get_conversation_by_id("old")).content == original + await db.engine.dispose() + + +@pytest.mark.asyncio +async def test_native_append_and_legacy_nested_edits_share_runtime(database): + from astrbot.core.agent.message import Message, dump_messages_with_checkpoints + from astrbot.core.agent.run_context import ContextWrapper + + conv = await database.create_conversation("umo", "p") + writer = ConversationEventWriter(database.conversation_store, await database.conversation_store.read(conv.conversation_id)) + await writer.start_turn({"kind": "agent"}) + writer.runtime_context = ContextWrapper(context=None, messages=[Message(role="user", content="input")]) + writer.stage_history([{"role": "user", "content": "input"}]) + native = await writer.append_message({"role": "assistant", "content": [{"type": "text", "text": "native"}]}, event_id="native") + assert len(writer.runtime_context.messages) == 2 + await writer.append_message({"role": "user", "content": "journal only"}, include_in_context=False) + assert writer.runtime_context.messages[-1]._no_save + writer.runtime_context.messages[1].content[0].text = "plugin edited" + history = dump_messages_with_checkpoints([m for m in writer.runtime_context.messages if not m._no_save]) + await writer.save_history(history) + restored = await database.conversation_store.read(conv.conversation_id) + assert restored.messages == history + assert len(restored.messages) == 2 + assert restored.messages[1]["content"][0]["text"] == "plugin edited" + events = await database.conversation_store.events(conv.conversation_id) + assert next(e for e in events if e.event_id == native.event_id).payload["message"]["content"][0]["text"] == "native" + + +@pytest.mark.asyncio +async def test_large_baseline_does_not_repeat_snapshots(database): + store = database.conversation_store + conv = await database.create_conversation("umo", "p", [{"role": "user", "content": "x" * (4 * 1024 * 1024 + 1)}]) + before = await store.read(conv.conversation_id) + events = await store.append(conv.conversation_id, [{"type": "message.appended", "payload": {"message": {"role": "assistant", "content": "small"}}}]) + assert [e.type for e in events] == ["message.appended"] + after = await store.read(conv.conversation_id) + assert after.conversation.replay_from_event_id == before.conversation.replay_from_event_id + assert after.replay_bytes < 1024 + + +@pytest.mark.asyncio +async def test_execution_protocol_rejects_invalid_relationships_and_usage(database): + conv = await database.create_conversation("umo", "p") + store = database.conversation_store + await store.append(conv.conversation_id, [{"event_id": "turn", "type": "turn.started", "payload": {}}]) + await store.append(conv.conversation_id, [{"event_id": "request", "type": "request.started", "payload": {"turn_id": "turn"}}]) + with pytest.raises(ValueError, match="subset"): + await store.append(conv.conversation_id, [{"type": "request.finished", "payload": {"request_id": "request", "status": "completed", "usage": {"input_tokens": 1, "cached_input_tokens": 2}}}]) + await store.append(conv.conversation_id, [{"type": "request.finished", "payload": {"request_id": "request", "status": "failed"}}]) + with pytest.raises(ConversationConflictError, match="finished"): + await store.append(conv.conversation_id, [{"type": "request.finished", "payload": {"request_id": "request", "status": "completed"}}]) + with pytest.raises(ValueError, match="parent turn"): + await store.append(conv.conversation_id, [{"type": "tool.started", "payload": {"turn_id": "request"}}]) + + +@pytest.mark.asyncio +async def test_old_backup_import_and_new_backup_round_trip(database): + from astrbot.core.backup.exporter import AstrBotExporter + from astrbot.core.backup.importer import AstrBotImporter + + importer = AstrBotImporter(database) + original = [{"role": "user", "content": "backup"}, {"role": "assistant", "content": "restored"}, {"role": "_checkpoint", "content": {"id": "cp"}}] + counts = await importer._import_main_database({ + "conversations": [{"inner_conversation_id": 7, "conversation_id": "old-backup", "platform_id": "webchat", "user_id": "webchat:FriendMessage:webchat!alice!s1", "content": original, "token_usage": 42}], + "platform_message_history": [{"id": 9, "platform_id": "webchat", "user_id": "s1", "content": {"type": "bot", "message": []}, "llm_checkpoint_id": "cp"}], + }) + assert counts["conversations"] == 1 + restored = await database.get_conversation_by_id("old-backup") + assert restored.content == original[:2] + linked = await database.get_platform_message_history_by_id(9) + assert linked.context_event_id and linked.turn_id != "cp" + exported = await AstrBotExporter(database)._export_main_database() + assert "conversations" not in exported and exported["conversation_events"] + await importer._clear_main_db() + await importer._import_main_database(exported) + assert (await database.get_conversation_by_id("old-backup")).content == original[:2] + assert (await database.get_platform_message_history_by_id(9)).context_event_id == linked.context_event_id + + +@pytest.mark.asyncio +async def test_temporary_parts_are_recorded_once_and_excluded_from_projection(database): + from astrbot.core.agent.message import Message, TextPart, dump_messages_with_checkpoints + conv = await database.create_conversation("umo", "p") + writer = ConversationEventWriter(database.conversation_store, await database.conversation_store.read(conv.conversation_id)) + await writer.start_turn({"kind": "agent"}) + messages = [Message(role="user", content=[TextPart(text="keep"), TextPart(text="temporary").mark_as_temp()])] + raw = dump_messages_with_checkpoints(messages, include_temporary=True) + writer.stage_history(raw) + writer.stage_history(raw) + await writer.save_history(raw) + events = await database.conversation_store.events(conv.conversation_id) + assert len([e for e in events if e.payload.get("include_in_context") is False]) == 1 + assert "temporary" in json.dumps([e.payload for e in events]) + projected = (await database.conversation_store.read(conv.conversation_id)).messages + assert projected[0]["content"] == [{"type": "text", "text": "keep"}] + + +@pytest.mark.asyncio +async def test_provider_retry_has_distinct_attempts(database): + from astrbot.core.agent.event_stream import RequestEventRecorder + from astrbot.core.provider.entities import TokenUsage + from astrbot.core.provider.sources.request_retry import retry_provider_request + + conv = await database.create_conversation("umo", "p") + writer = ConversationEventWriter(database.conversation_store, await database.conversation_store.read(conv.conversation_id)) + await writer.start_turn({"kind": "agent"}) + recorder = RequestEventRecorder(writer.consume, {"turn_id": writer.turn_id}) + await recorder.begin() + attempts = 0 + + class RetryableError(RuntimeError): + status_code = 429 + + async def invoke(): + nonlocal attempts + attempts += 1 + if attempts == 1: + raise RetryableError("rate limited") + return "answer" + + assert await retry_provider_request( + "test", invoke, max_attempts=2, request_event_recorder=recorder + ) == "answer" + await recorder.finish("completed", usage=TokenUsage(input_other=10, output=3)) + events = await database.conversation_store.events(conv.conversation_id) + starts = [e for e in events if e.type == "request.started"] + finishes = [e for e in events if e.type == "request.finished"] + assert len(starts) == len(finishes) == 2 + assert [e.payload["request_id"] for e in finishes] == [e.event_id for e in starts] + assert [e.payload["status"] for e in finishes] == ["failed", "completed"] + + +@pytest.mark.asyncio +async def test_idempotent_batch_includes_its_automatic_rebase(database): + conv = await database.create_conversation("umo", "p") + store = database.conversation_store + drafts = [{"event_id": f"m{i}", "type": "message.appended", "payload": {"message": {"role": "user", "content": str(i)}}} for i in range(257)] + first = await store.append(conv.conversation_id, drafts, expected_head=1, expected_leaf=None) + assert first[-1].payload["reason"] == "snapshot" + retried = await store.append(conv.conversation_id, drafts, expected_head=1, expected_leaf=None) + assert [e.event_id for e in retried] == [e.event_id for e in first] + assert (await store.read(conv.conversation_id)).replay_count == 1 + + +@pytest.mark.asyncio +async def test_webchat_links_plugin_replacement_without_an_appended_reply(database): + store = database.conversation_store + conv = await database.create_conversation("umo", "p", [{"role": "user", "content": "old"}]) + writer = ConversationEventWriter(store, await store.read(conv.conversation_id)) + await writer.start_turn({"kind": "agent"}) + await writer.save_history([{"role": "user", "content": "rewritten"}, {"role": "assistant", "content": "new answer"}]) + record = await database.insert_platform_message_history("webchat", "s1", {"type": "bot", "message": []}, turn_id=writer.turn_id) + assert record.context_event_id == writer.leaf_event_id + fork = await store.create(umo="umo", platform_id="p", parent_event_id=record.context_event_id) + assert (await store.read(fork.conversation_id)).messages[-1]["content"] == "new answer" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("change", ["append", "select_branch"]) +async def test_declared_read_revision_rejects_context_changes(database, change): + from dataclasses import FrozenInstanceError + + from astrbot.core.conversation_mgr import ConversationManager + from astrbot.core.db.po import ConversationRead, ConversationRevision + + store = database.conversation_store + manager = ConversationManager(database) + created = await database.create_conversation("umo", "p") + first = (await store.append(created.conversation_id, [ + {"type": "message.appended", "payload": {"message": {"role": "user", "content": "first"}}}, + {"type": "message.appended", "payload": {"message": {"role": "assistant", "content": "second"}}}, + ]))[0] + record = await database.get_conversation_by_id(created.conversation_id) + conversation = await manager.get_conversation("umo", created.conversation_id) + assert isinstance(record, ConversationRead) + assert isinstance(record.revision, ConversationRevision) + assert conversation.revision == record.revision + assert json.loads(conversation.history) == record.content + assert "revision" not in ConversationV3.__table__.columns + assert "revision" not in ConversationV2.__table__.columns + with pytest.raises(FrozenInstanceError): + conversation.revision.head_seq = 0 + writer = await manager.event_writer("umo", created.conversation_id, expected_revision=conversation.revision) + if change == "append": + await writer.append("plugin.test.changed", {}) + else: + await store.select_branch( + created.conversation_id, first.event_id, + expected_head=record.revision.head_seq, + expected_leaf=record.revision.leaf_event_id, + ) + with pytest.raises(ConversationConflictError, match="preparing"): + await manager.event_writer("umo", created.conversation_id, expected_revision=conversation.revision) + + +@pytest.mark.asyncio +async def test_autonomous_summary_uses_the_same_writer_after_execution_events(database): + from types import SimpleNamespace + + from astrbot.core.agent.response import AgentResponse + from astrbot.core.provider.entities import ProviderRequest + from astrbot.core.utils.history_saver import persist_agent_history + + conv = await database.create_conversation("umo", "p") + store = database.conversation_store + writer = ConversationEventWriter(store, await store.read(conv.conversation_id)) + await writer.consume(AgentResponse("turn.started", {}, "summary-turn")) + await writer.consume(AgentResponse("request.started", {}, "summary-request")) + await writer.consume(AgentResponse("request.finished", { + "request_id": "summary-request", "status": "completed", + })) + req = ProviderRequest() + req.conversation = SimpleNamespace(history="[]", cid=conv.conversation_id) + await persist_agent_history( + None, event=SimpleNamespace(unified_msg_origin="umo"), req=req, + summary_note="Task completed", conversation_events=writer, + ) + await writer.finish_turn("completed") + assert (await store.read(conv.conversation_id)).messages == [ + {"role": "user", "content": "Output your last task result below."}, + {"role": "assistant", "content": "Task completed"}, + ] + events = await store.events(conv.conversation_id) + assert [e.type for e in events] == [ + "conversation.created", "turn.started", "request.started", "request.finished", + "message.appended", "message.appended", "turn.finished", + ] + + +@pytest.mark.asyncio +async def test_display_context_lookup_holds_the_write_transaction(database, monkeypatch): + """An event commit must not slip between link lookup and display insertion.""" + import sqlite3 + + from sqlalchemy.ext.asyncio import AsyncSession + from sqlalchemy.sql import Select + + store = database.conversation_store + conv = await store.create(umo="umo", platform_id="webchat") + writer = ConversationEventWriter(store, await store.read(conv.conversation_id)) + await writer.start_turn({"kind": "agent"}) + original_execute = AsyncSession.execute + checked = False + + async def execute(session, statement, *args, **kwargs): + nonlocal checked + result = await original_execute(session, statement, *args, **kwargs) + if isinstance(statement, Select) and any( + column.get("entity") is ConversationEvent + for column in statement.column_descriptions + ): + contender = sqlite3.connect(database.engine.url.database, timeout=0) + try: + with pytest.raises(sqlite3.OperationalError, match="locked"): + contender.execute("BEGIN IMMEDIATE") + checked = True + finally: + contender.close() + return result + + with monkeypatch.context() as patch: + patch.setattr(AsyncSession, "execute", execute) + record = await database.insert_platform_message_history( + "webchat", "session", {"type": "bot", "message": []}, + turn_id=writer.turn_id, + ) + assert checked + assert record.context_event_id is None + await writer.save_history([{"role": "assistant", "content": "answer"}]) + linked = await database.get_platform_message_history_by_id(record.id) + assert linked.context_event_id == writer.leaf_event_id + + +@pytest.mark.asyncio +@pytest.mark.parametrize("display_first", [False, True]) +async def test_edited_turn_reply_can_fork_without_original_or_later_messages(database, display_first): + store = database.conversation_store + umo = "webchat:FriendMessage:webchat!alice!edited-session" + conv = await store.create(umo=umo, platform_id="webchat") + writer = ConversationEventWriter(store, await store.read(conv.conversation_id)) + await writer.start_turn({"kind": "agent"}) + user = await database.insert_platform_message_history( + "webchat", "edited-session", {"type": "user", "message": []}, turn_id=writer.turn_id, + ) + await writer.save_history([ + {"role": "user", "content": "original"}, + {"role": "assistant", "content": "original answer"}, + ]) + await writer.finish_turn("completed") + snapshot = await store.read(conv.conversation_id) + replacement = await store.rewind_webchat( + conv.conversation_id, user.id, + content={"type": "user", "message": [{"type": "plain", "text": "edited"}]}, + expected_head=snapshot.conversation.head_seq, + expected_leaf=snapshot.conversation.leaf_event_id, + ) + writer = ConversationEventWriter(store, await store.read(conv.conversation_id)) + await writer.start_turn({"kind": "agent"}, event_id=replacement.turn_id) + messages = [{"role": "user", "content": "edited"}, {"role": "assistant", "content": "edited answer"}] + if not display_first: + await writer.save_history(messages) + bot = await database.insert_platform_message_history( + "webchat", "edited-session", {"type": "bot", "message": []}, turn_id=replacement.turn_id, + ) + if display_first: + await writer.save_history(messages) + await writer.finish_turn("completed") + linked = await database.get_platform_message_history_by_id(bot.id) + assert linked.context_event_id + later = ConversationEventWriter(store, await store.read(conv.conversation_id)) + await later.append_message({"role": "user", "content": "later message"}) + branch = await store.create(umo=umo, platform_id="webchat", parent_event_id=linked.context_event_id) + assert (await store.read(branch.conversation_id)).messages == messages diff --git a/tests/test_tool_loop_agent_runner.py b/tests/test_tool_loop_agent_runner.py index 180e0edf2d..be1fd46a8f 100644 --- a/tests/test_tool_loop_agent_runner.py +++ b/tests/test_tool_loop_agent_runner.py @@ -840,13 +840,13 @@ async def test_empty_messages_after_on_agent_begin_skip_provider( assert runner.done() assert not runner.was_aborted() assert runner.run_context.messages == [] - assert responses[-1].type == "err" + assert [r for r in responses if "chain" in r.data][-1].type == "err" final_response = runner.get_final_llm_resp() assert final_response is not None assert final_response.role == "err" assert final_response.completion_text == "No messages remain for the LLM request." assert ( - responses[-1].data["chain"].get_plain_text() + [r for r in responses if "chain" in r.data][-1].data["chain"].get_plain_text() == "LLM 响应错误: No messages remain for the LLM request." ) @@ -871,7 +871,7 @@ async def test_empty_request_after_on_llm_request_skip_provider( assert runner.done() assert not runner.was_aborted() assert runner.run_context.messages == [] - assert responses[-1].type == "err" + assert [r for r in responses if "chain" in r.data][-1].type == "err" final_response = runner.get_final_llm_resp() assert final_response is not None assert final_response.role == "err" @@ -1371,6 +1371,23 @@ async def test_empty_output_retries_exhausted_then_uses_fallback_provider( assert fallback_provider.call_count == 1 +async def display_responses(source): + """Consume protocol events while exposing display responses to UI tests. + + Args: + source: Runner response stream. + + Yields: + Responses carrying display message chains. + """ + from contextlib import aclosing + + async with aclosing(source): + async for response in source: + if "chain" in response.data: + yield response + + @pytest.mark.asyncio async def test_stop_signal_returns_aborted_and_discards_partial_message( runner, provider_request, mock_tool_executor, mock_hooks @@ -1386,7 +1403,7 @@ async def test_stop_signal_returns_aborted_and_discards_partial_message( streaming=True, ) - step_iter = runner.step() + step_iter = display_responses(runner.step()) first_resp = await step_iter.__anext__() assert first_resp.type == "streaming_delta" @@ -1439,7 +1456,7 @@ async def test_stop_cancels_provider_before_first_response( streaming=streaming, ) - step_iter = runner.step() + step_iter = display_responses(runner.step()) pending_response = asyncio.create_task(anext(step_iter)) await asyncio.wait_for(provider.started.wait(), timeout=1) @@ -1495,7 +1512,7 @@ async def test_stop_interrupts_pending_subagent_handoff(mock_hooks): streaming=False, ) - step_iter = runner.step() + step_iter = display_responses(runner.step()) first_resp = await step_iter.__anext__() if first_resp.type == "agent_stats": first_resp = await step_iter.__anext__() @@ -1554,7 +1571,7 @@ async def test_stop_interrupts_pending_regular_tool(mock_hooks): streaming=False, ) - step_iter = runner.step() + step_iter = display_responses(runner.step()) first_resp = await step_iter.__anext__() if first_resp.type == "agent_stats": first_resp = await step_iter.__anext__() @@ -1803,8 +1820,8 @@ async def text_chat(self, **kwargs) -> LLMResponse: async def recorded_step(): async for response in original_step(): - chain = response.data["chain"] - if chain.get_plain_text() in (final_text, reasoning): + chain = response.data.get("chain") + if chain and chain.get_plain_text() in (final_text, reasoning): final_events.append((response.type, chain.type)) hooks_at_emission.append(mock_hooks.agent_done_called) yield response @@ -2274,3 +2291,187 @@ async def test_follow_up_after_stop_not_merged_into_tool_result( if __name__ == "__main__": # 运行测试 pytest.main([__file__, "-v"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True]) +async def test_event_journal_records_requests_tools_and_final_context( + tmp_path, runner, provider_request, mock_tool_executor, mock_hooks, streaming +): + from astrbot.core.agent.conversation_events import ConversationEventWriter + from astrbot.core.agent.message import dump_messages_with_checkpoints + from astrbot.core.db.sqlite import SQLiteDatabase + + db = SQLiteDatabase(str(tmp_path / "runner_events.db")) + await db.initialize() + try: + conv = await db.create_conversation("umo", "p") + writer = ConversationEventWriter(db.conversation_store, await db.conversation_store.read(conv.conversation_id)) + provider = VaryingUsageProvider() + await runner.reset( + provider=provider, request=provider_request, + run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, agent_hooks=mock_hooks, + streaming=streaming, + ) + writer.runtime_context = runner.run_context + async for response in runner.step_until_done(3): + await writer.consume(response) + await writer.save_history(dump_messages_with_checkpoints(runner.run_context.messages)) + await writer.finish_turn("completed") + events = await db.conversation_store.events(conv.conversation_id) + types = [e.type for e in events] + assert types.count("request.started") == types.count("request.finished") == 2 + assert types.count("tool.started") == types.count("tool.finished") == 1 + assert types[-1] == "turn.finished" + usage = [e.payload["usage"] for e in events if e.type == "request.finished"] + assert sum(item["input_tokens"] for item in usage) == 330 + assert sum(item["cached_input_tokens"] for item in usage) == 30 + finished_tool = next(e for e in events if e.type == "tool.finished") + assert finished_tool.payload["status"] == "completed" + assert finished_tool.payload["result"]["content"][0]["type"] == "text" + history = (await db.conversation_store.read(conv.conversation_id)).messages + assert [m["role"] for m in history] == ["user", "assistant", "tool", "assistant"] + assert history[-1]["content"] == [{"type": "text", "text": "final"}] + finally: + await db.engine.dispose() + + +@pytest.mark.asyncio +async def test_journal_failure_does_not_retry_a_model_or_fallback( + runner, provider_request, mock_tool_executor, mock_hooks +): + from astrbot.core.agent.conversation_events import ConversationPersistenceError + + primary, fallback = MockProvider(), MockProvider() + await runner.reset( + provider=primary, fallback_providers=[fallback], request=provider_request, + run_context=ContextWrapper(context=None), tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, streaming=False, + ) + from contextlib import aclosing + + async with aclosing(runner.step()) as responses: + with pytest.raises(ConversationPersistenceError): + async for response in responses: + if response.type == "request.started": + raise ConversationPersistenceError("unavailable") + assert primary.call_count == fallback.call_count == 0 + + +@pytest.mark.asyncio +async def test_tool_started_is_acknowledged_before_creating_executor( + runner, provider_request, mock_hooks +): + from contextlib import aclosing + from unittest.mock import Mock + + provider = VaryingUsageProvider() + executor = SimpleNamespace(execute=Mock(side_effect=AssertionError("must not run"))) + await runner.reset( + provider=provider, request=provider_request, + run_context=ContextWrapper(context=None), tool_executor=executor, + agent_hooks=mock_hooks, streaming=False, + ) + seen = [] + async with aclosing(runner.step_until_done(3)) as responses: + async for response in responses: + seen.append(response.type) + if response.type == "request.started": + assert provider.call_count == 0 + if response.type == "tool.started": + assert response.event_id + assert response.data["turn_id"] == runner.turn_id + executor.execute.assert_not_called() + break + executor.execute.assert_not_called() + assert seen.index("request.finished") < seen.index("tool.started") + assert not runner._events.active + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streaming", [False, True]) +async def test_cancelled_consumer_joins_provider_work( + runner, provider_request, mock_tool_executor, mock_hooks, streaming +): + provider = MockBlockingProvider() + await runner.reset( + provider=provider, request=provider_request, + run_context=ContextWrapper(context=None), tool_executor=mock_tool_executor, + agent_hooks=mock_hooks, streaming=streaming, + ) + source = runner.step_until_done(3) + async for response in source: + if response.type == "request.started": + break + pending = asyncio.create_task(anext(source)) + await asyncio.wait_for(provider.started.wait(), 1) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, 1) + await source.aclose() + assert provider.cancelled.is_set() + assert not runner._events.active + + +@pytest.mark.asyncio +async def test_summary_requests_share_the_runner_event_stream( + runner, mock_tool_executor, mock_hooks +): + primary = MockProvider() + primary.provider_config["max_context_tokens"] = 100 + summary = MockProvider() + request = ProviderRequest( + prompt="Continue", contexts=[ + {"role": "user", "content": "Long history " * 300}, + {"role": "assistant", "content": "Previous response " * 300}, + ], + ) + await runner.reset( + provider=primary, request=request, run_context=ContextWrapper(context=None), + tool_executor=mock_tool_executor, agent_hooks=mock_hooks, streaming=False, + llm_compress_provider=summary, + ) + events = [response async for response in runner.step_until_done(1)] + starts = [response for response in events if response.type == "request.started"] + finishes = [response for response in events if response.type == "request.finished"] + assert summary.call_count == primary.call_count == 1 + assert len(starts) == len(finishes) == 2 + assert starts[0].data["purpose"] == "compaction" + assert [e.data["request_id"] for e in finishes] == [e.event_id for e in starts] + assert all(e.data["turn_id"] == runner.turn_id for e in starts) + + +@pytest.mark.asyncio +async def test_cancelled_consumer_joins_pending_tool( + runner, provider_request, mock_hooks +): + started, cancelled = asyncio.Event(), asyncio.Event() + + class BlockingExecutor: + async def execute(self, **kwargs): + started.set() + try: + await asyncio.Future() + finally: + cancelled.set() + if False: + yield None + + await runner.reset( + provider=VaryingUsageProvider(), request=provider_request, + run_context=ContextWrapper(context=None), tool_executor=BlockingExecutor(), + agent_hooks=mock_hooks, streaming=False, + ) + source = runner.step_until_done(3) + async for response in source: + if response.type == "tool.started": + break + pending = asyncio.create_task(anext(source)) + await asyncio.wait_for(started.wait(), 1) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(pending, 1) + await source.aclose() + assert cancelled.is_set() + assert not runner._events.active diff --git a/tests/unit/test_astr_agent_tool_exec.py b/tests/unit/test_astr_agent_tool_exec.py index 675367dcec..913beb0f03 100644 --- a/tests/unit/test_astr_agent_tool_exec.py +++ b/tests/unit/test_astr_agent_tool_exec.py @@ -390,13 +390,24 @@ async def test_background_wakeup_passes_history_and_provider_settings_to_main_ag {"role": "assistant", "content": "old answer"}, ] captured: dict = {} + writer = SimpleNamespace(consume=AsyncMock(), finish_turn=AsyncMock()) + from astrbot.core.agent.response import AgentResponse + + started = AgentResponse("turn.started", {}, "turn") + + class EventRunner(_DoneRunner): + run_context = SimpleNamespace(messages=[]) + + async def step_until_done(self, _max_step): + yield started + async def _fake_get_session_conv(**_kwargs): return SimpleNamespace(history=json.dumps(history)) async def _fake_build_main_agent(**kwargs): captured.update(kwargs) - return SimpleNamespace(agent_runner=_DoneRunner()) + return SimpleNamespace(agent_runner=EventRunner(), conversation_events=writer) monkeypatch.setattr( "astrbot.core.astr_main_agent._get_session_conv", @@ -448,6 +459,11 @@ async def _fake_build_main_agent(**kwargs): assert "old answer" not in request.system_prompt assert request.contexts == history + writer.consume.assert_awaited_once_with(started) + writer.finish_turn.assert_awaited_once_with("completed") + assert writer.runtime_context is EventRunner.run_context + assert captured["event"].conversation_events is None + @pytest.mark.asyncio @pytest.mark.parametrize( @@ -483,7 +499,7 @@ async def _fake_get_session_conv(**_kwargs): return SimpleNamespace(history="[]") async def _fake_build_main_agent(**_kwargs): - return SimpleNamespace(agent_runner=runner) + return SimpleNamespace(agent_runner=runner, conversation_events=None) monkeypatch.setattr( "astrbot.core.astr_main_agent._get_session_conv", diff --git a/tests/unit/test_astr_main_agent.py b/tests/unit/test_astr_main_agent.py index e41380fcde..c18d460dc9 100644 --- a/tests/unit/test_astr_main_agent.py +++ b/tests/unit/test_astr_main_agent.py @@ -312,6 +312,7 @@ async def test_proactive_agent_respects_runtime_and_safety_settings( runner = runner_cls.return_value runner.reset = AsyncMock() runner.step_until_done.return_value.__aiter__.return_value = [] + runner.step_until_done.return_value.aclose = AsyncMock() runner.get_final_llm_resp.return_value = None if entrypoint == "cron": diff --git a/tests/unit/test_cron_context_compression.py b/tests/unit/test_cron_context_compression.py index c9d585dd59..10fb7727f7 100644 --- a/tests/unit/test_cron_context_compression.py +++ b/tests/unit/test_cron_context_compression.py @@ -145,7 +145,7 @@ async def steps(max_step): ), patch( "astrbot.core.astr_main_agent.build_main_agent", - AsyncMock(return_value=SimpleNamespace(agent_runner=runner)), + AsyncMock(return_value=SimpleNamespace(agent_runner=runner, conversation_events=None)), ) as build, patch("astrbot.core.cron.manager.persist_agent_history", AsyncMock()), patch("astrbot.core.astr_agent_tool_exec.persist_agent_history", AsyncMock()), diff --git a/tests/unit/test_cron_manager.py b/tests/unit/test_cron_manager.py index 0dcb480d77..eab07d8fd5 100644 --- a/tests/unit/test_cron_manager.py +++ b/tests/unit/test_cron_manager.py @@ -579,13 +579,19 @@ async def test_woke_main_agent_passes_history_and_provider_settings( conv = MagicMock() conv.history = json.dumps(history) + from astrbot.core.agent.response import AgentResponse + + started = AgentResponse("turn.started", {}, "turn") + writer = MagicMock(consume=AsyncMock(), finish_turn=AsyncMock()) + class FakeRunner: + run_context = MagicMock() + state = AgentState.DONE def step_until_done(self, max_step): async def gen(): - if False: - yield None + yield started return gen() @@ -597,7 +603,7 @@ def get_final_llm_resp(self): async def fake_build_main_agent(*, event, plugin_context, config, req): captured["config"] = config captured["req"] = req - return MagicMock(agent_runner=FakeRunner()) + return MagicMock(agent_runner=FakeRunner(), conversation_events=writer) async def fake_persist_agent_history(*args, **kwargs): return None @@ -630,6 +636,9 @@ async def fake_persist_agent_history(*args, **kwargs): assert "old question" not in request.system_prompt assert "old answer" not in request.system_prompt assert request.contexts == history + writer.consume.assert_awaited_once_with(started) + writer.finish_turn.assert_awaited_once_with("failed") + assert writer.runtime_context is FakeRunner.run_context @pytest.mark.asyncio @pytest.mark.parametrize( @@ -685,7 +694,7 @@ def get_final_llm_resp(self): runner = _StepCapturingRunner() async def fake_build_main_agent(*, event, plugin_context, config, req): - return MagicMock(agent_runner=runner) + return MagicMock(agent_runner=runner, conversation_events=None) with ( patch( @@ -739,7 +748,7 @@ def get_final_llm_resp(self): resp.completion_text = "malformed_function_call" return resp - fake_result = SimpleNamespace(agent_runner=FakeRunner()) + fake_result = SimpleNamespace(agent_runner=FakeRunner(), conversation_events=None) with ( patch( "astrbot.core.astr_main_agent.build_main_agent",