xiaoski/PicoBot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(session): commit same-turn media with final reply

Browse Source
This commit is contained in:
xiaoxixi 2026-07-20 00:11:39 +08:00
parent 4175af113c
commit b043fcf531
6 changed files with 397 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
src/session/messenger.rs
View File

@ -2,10 +2,13 @@ use std::collections::HashMap;
use crate::bus::{ChatMessage, MediaItem, MessageSource, OutboundMessage, SourceKind};
use crate::session::UnifiedSessionId;
use crate::tools::OutboundMessenger;
use crate::tools::{OutboundDelivery, OutboundMessenger};
use super::persistence::append_persisted_messages;
use super::session::{CURRENT_SOURCE_SESSION, SessionManager};
use super::persistence::{append_active_turn_message, append_persisted_messages};
use super::session::{
CURRENT_SOURCE_SESSION, CURRENT_TURN_DELIVERIES, CURRENT_TURN_ID, PendingTurnDelivery,
SessionManager,
};
#[async_trait::async_trait]
impl OutboundMessenger for SessionManager {
@ -17,7 +20,7 @@ impl OutboundMessenger for SessionManager {
content: &str,
mut source: MessageSource,
media: Vec<MediaItem>,
) -> Result<(), String> {
) -> Result<OutboundDelivery, String> {
if source.from_session.is_none() {
source.from_session = CURRENT_SOURCE_SESSION
.try_with(|value| value.clone())
@ -47,6 +50,32 @@ impl OutboundMessenger for SessionManager {
let origin = source.from_session.as_deref().unwrap_or("unknown");
let origin_id = source.from_session.clone();
let same_session = source.from_session.as_deref() == Some(target_sid.to_string().as_str());
let current_turn_id = CURRENT_TURN_ID
.try_with(|value| value.clone())
.ok()
.flatten();
let current_turn_deliveries = CURRENT_TURN_DELIVERIES
.try_with(|value| value.clone())
.ok()
.flatten();
if same_session
&& channel == "cli_chat"
&& !media.is_empty()
&& let (Some(turn_id), Some(deliveries)) =
(current_turn_id.clone(), current_turn_deliveries)
{
let guard = session.lock().await;
if guard.owns_active_turn(&turn_id) {
deliveries
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
.push(PendingTurnDelivery {
content: content.to_string(),
media,
});
return Ok(OutboundDelivery::AttachedToCurrentTurn);
}
}
let marked_content = if content.trim().is_empty() && !media.is_empty() && same_session {
String::new()
} else {
@ -55,9 +84,14 @@ impl OutboundMessenger for SessionManager {
let message = outbound_history_message(marked_content.clone(), source, &media);
let message_id = message.id.clone();
append_persisted_messages(&session, vec![message])
.await
.map_err(|error| error.to_string())?;
if same_session && let Some(turn_id) = current_turn_id {
// Ownership is revalidated atomically with the in-memory append;
// a concurrent /stop therefore falls back to a versioned write.
append_active_turn_message(&session, message, turn_id).await
} else {
append_persisted_messages(&session, vec![message]).await
}
.map_err(|error| error.to_string())?;
if let Some(origin_id) = origin_id {
self.restore_origin_dialog(&origin_id, &target_sid).await;
@ -78,7 +112,8 @@ impl OutboundMessenger for SessionManager {
delivery: None,
})
.await
.map_err(|error| error.to_string())
.map_err(|error| error.to_string())?;
Ok(OutboundDelivery::Delivered)
}
}
@ -108,6 +143,7 @@ impl SessionManager {
Vec::new(),
)
.await
.map(|_| ())
}
}

129
src/session/persistence.rs
View File

@ -49,9 +49,42 @@ pub(super) async fn append_persisted_messages(
.map(|_| ())
}
enum VersionPolicy {
Advance,
PreserveForOwnedTurn(String),
}
/// Persist a message emitted as a side effect of the currently active Turn
/// without invalidating that same Turn's captured session version.
///
/// The task-local Turn ID is revalidated while holding the session lock. If it
/// no longer owns the target session, the write advances `state_version` like
/// any unrelated mutation so stale model work is still rejected.
pub(super) async fn append_active_turn_message(
session: &Arc<Mutex<Session>>,
message: ChatMessage,
turn_id: String,
) -> Result<(), StorageError> {
append_persisted_messages_inner(
session,
vec![message],
VersionPolicy::PreserveForOwnedTurn(turn_id),
)
.await
.map(|_| ())
}
pub(super) async fn append_persisted_messages_with_meta(
session: &Arc<Mutex<Session>>,
messages: Vec<ChatMessage>,
) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> {
append_persisted_messages_inner(session, messages, VersionPolicy::Advance).await
}
async fn append_persisted_messages_inner(
session: &Arc<Mutex<Session>>,
messages: Vec<ChatMessage>,
version_policy: VersionPolicy,
) -> Result<Vec<crate::storage::message::MessageMeta>, StorageError> {
if messages.is_empty() {
return Ok(Vec::new());
@ -60,12 +93,19 @@ pub(super) async fn append_persisted_messages_with_meta(
let persistence_lock = { session.lock().await.persistence_lock.clone() };
let _persistence_guard = persistence_lock.lock().await;
let message_ids: Vec<_> = messages.iter().map(|message| message.id.clone()).collect();
let snapshots: Vec<Option<MessagePersistSnapshot>> = {
let (snapshots, advance_state_version): (Vec<Option<MessagePersistSnapshot>>, bool) = {
let mut guard = session.lock().await;
messages
let advance_state_version = match &version_policy {
VersionPolicy::Advance => true,
VersionPolicy::PreserveForOwnedTurn(turn_id) => !guard.owns_active_turn(turn_id),
};
let snapshots = messages
.into_iter()
.map(|message| guard.add_message_in_memory(message, true))
.collect()
.map(|message| {
guard.add_message_in_memory_with_version(message, true, advance_state_version)
})
.collect();
(snapshots, advance_state_version)
};
let committed = snapshots
.iter()
@ -74,7 +114,10 @@ pub(super) async fn append_persisted_messages_with_meta(
.collect();
if let Err(error) = persist_added_messages(snapshots).await {
session.lock().await.rollback_message_suffix(&message_ids);
session
.lock()
.await
.rollback_message_suffix_with_version(&message_ids, advance_state_version);
return Err(error);
}
Ok(committed)
@ -109,7 +152,12 @@ where
#[cfg(test)]
mod tests {
use super::*;
use crate::config::LLMProviderConfig;
use crate::memory::MemoryManager;
use crate::session::{TurnController, TurnStatus};
use crate::tools::ToolRegistry;
use std::collections::HashMap;
use std::path::PathBuf;
#[tokio::test]
async fn completed_is_published_only_after_persistence_succeeds() {
@ -149,4 +197,75 @@ mod tests {
Some("failed to persist turn: database down")
);
}
#[tokio::test]
async fn active_turn_side_effect_does_not_invalidate_its_session_version() {
let dir = tempfile::tempdir().unwrap();
let storage = Arc::new(
crate::storage::Storage::new(&dir.path().join("memory.db"))
.await
.unwrap(),
);
let memory_manager = Arc::new(MemoryManager::new(
storage,
"test".to_string(),
"test".to_string(),
));
let config = LLMProviderConfig {
provider_type: "openai".to_string(),
name: "test".to_string(),
base_url: "http://127.0.0.1".to_string(),
api_key: "test".to_string(),
extra_headers: HashMap::new(),
model_id: "test".to_string(),
temperature: None,
max_tokens: None,
model_extra: HashMap::new(),
max_tool_iterations: 1,
token_limit: 8_192,
workspace_dir: PathBuf::from("."),
input_types: vec!["text".to_string(), "image".to_string()],
};
let session = Arc::new(Mutex::new(
Session::new(
crate::session::UnifiedSessionId::new("cli_chat", "chat", "dialog"),
config,
Arc::new(ToolRegistry::new()),
None,
String::new(),
"test".to_string(),
memory_manager,
)
.await
.unwrap(),
));
let base_version = session.lock().await.state_version_for_test();
session.lock().await.set_active_turn_for_test("turn-1");
append_active_turn_message(
&session,
ChatMessage::assistant("sent screenshot"),
"turn-1".to_string(),
)
.await
.unwrap();
{
let guard = session.lock().await;
assert_eq!(guard.state_version_for_test(), base_version);
assert_eq!(guard.get_history().len(), 1);
}
append_active_turn_message(
&session,
ChatMessage::assistant("late screenshot"),
"stale-turn".to_string(),
)
.await
.unwrap();
let guard = session.lock().await;
assert_eq!(guard.state_version_for_test(), base_version + 1);
assert_eq!(guard.get_history().len(), 2);
}
}

221
src/session/session.rs
View File

@ -73,6 +73,117 @@ fn committed_turn_delta(
tokio::task_local! {
pub(super) static CURRENT_SOURCE_SESSION: Option<String>;
pub(super) static CURRENT_TURN_ID: Option<String>;
pub(super) static CURRENT_TURN_DELIVERIES: Option<Arc<std::sync::Mutex<Vec<PendingTurnDelivery>>>>;
}
pub(super) struct PendingTurnDelivery {
pub content: String,
pub media: Vec<MediaItem>,
}
fn take_pending_turn_deliveries(
deliveries: &Arc<std::sync::Mutex<Vec<PendingTurnDelivery>>>,
) -> Vec<PendingTurnDelivery> {
std::mem::take(
&mut *deliveries
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner()),
)
}
fn take_current_turn_deliveries() -> Vec<PendingTurnDelivery> {
CURRENT_TURN_DELIVERIES
.try_with(|deliveries| {
deliveries
.as_ref()
.map(take_pending_turn_deliveries)
.unwrap_or_default()
})
.unwrap_or_default()
}
fn collapse_pending_turn_deliveries(pending: Vec<PendingTurnDelivery>) -> (String, Vec<MediaRef>) {
let fallback_content = pending
.iter()
.map(|delivery| delivery.content.trim())
.filter(|content| !content.is_empty())
.collect::<Vec<_>>()
.join("\n\n");
let mut media_refs = Vec::new();
for media_ref in pending
.into_iter()
.flat_map(|delivery| delivery.media.into_iter().map(|media| media.to_media_ref()))
{
if !media_refs.iter().any(|existing: &MediaRef| {
existing.path == media_ref.path && existing.media_type == media_ref.media_type
}) {
media_refs.push(media_ref);
}
}
(fallback_content, media_refs)
}
fn attach_pending_to_message(
message: &mut ChatMessage,
fallback_content: &str,
media_refs: &[MediaRef],
) {
if message.content.trim().is_empty() && !fallback_content.is_empty() {
message.content = fallback_content.to_string();
}
for media_ref in media_refs {
if !message.media_refs.iter().any(|existing| {
existing.path == media_ref.path && existing.media_type == media_ref.media_type
}) {
message.media_refs.push(media_ref.clone());
}
}
}
fn attach_pending_turn_deliveries(
result: &mut crate::agent::AgentProcessResult,
pending: Vec<PendingTurnDelivery>,
) {
let (fallback_content, media_refs) = collapse_pending_turn_deliveries(pending);
if media_refs.is_empty() && fallback_content.is_empty() {
return;
}
attach_pending_to_message(&mut result.final_response, &fallback_content, &media_refs);
if let Some(final_message) = result
.emitted_messages
.iter_mut()
.rev()
.find(|message| message.id == result.final_response.id)
{
final_message
.content
.clone_from(&result.final_response.content);
final_message
.media_refs
.clone_from(&result.final_response.media_refs);
}
}
fn partial_assistant_with_pending_deliveries(
snapshot: &TurnSnapshot,
completion_status: CompletionStatus,
pending: Vec<PendingTurnDelivery>,
) -> Option<ChatMessage> {
let (fallback_content, media_refs) = collapse_pending_turn_deliveries(pending);
let mut message = partial_assistant_message(snapshot, completion_status).or_else(|| {
(!media_refs.is_empty()).then(|| {
let mut message = ChatMessage::assistant("");
message.id = snapshot.message_id.clone();
message.turn_id = Some(snapshot.id.0.clone());
message.completion_status = completion_status;
message
})
})?;
attach_pending_to_message(&mut message, &fallback_content, &media_refs);
Some(message)
}
/// Result of handling a message - either an AI response or a command output
@ -182,7 +293,11 @@ async fn fail_turn_with_partial(
error: String,
) {
let snapshot = controller.snapshot();
let partial = partial_assistant_message(&snapshot, CompletionStatus::Interrupted);
let partial = partial_assistant_with_pending_deliveries(
&snapshot,
CompletionStatus::Interrupted,
take_current_turn_deliveries(),
);
if let Some(partial) = partial {
controller.begin_finalizing();
if let Err(persistence_error) = append_persisted_messages(session, vec![partial]).await {
@ -369,6 +484,60 @@ mod cancelled_partial_tests {
.is_none()
);
}
#[test]
fn pending_same_turn_media_is_attached_to_the_final_response() {
let final_message = ChatMessage::assistant("截图已经准备好了");
let final_id = final_message.id.clone();
let mut result = crate::agent::AgentProcessResult {
final_response: final_message.clone(),
emitted_messages: vec![
ChatMessage::tool("call-1", "send_message", "附件已加入当前回复"),
final_message,
],
total_tokens: None,
usage: None,
};
attach_pending_turn_deliveries(
&mut result,
vec![PendingTurnDelivery {
content: "这是百度首页截图".to_string(),
media: vec![MediaItem::new("/tmp/baidu.png", "image")],
}],
);
assert_eq!(result.final_response.content, "截图已经准备好了");
assert_eq!(result.final_response.media_refs.len(), 1);
assert_eq!(result.final_response.media_refs[0].path, "/tmp/baidu.png");
let committed_final = result
.emitted_messages
.iter()
.find(|message| message.id == final_id)
.unwrap();
assert_eq!(committed_final.media_refs.len(), 1);
assert_eq!(committed_final.media_refs[0].path, "/tmp/baidu.png");
}
#[test]
fn pending_media_survives_a_turn_without_partial_text() {
let (controller, _emitter, _) = TurnController::start("session", "message-id");
let message = partial_assistant_with_pending_deliveries(
&controller.snapshot(),
CompletionStatus::Interrupted,
vec![PendingTurnDelivery {
content: "这是已生成的截图".to_string(),
media: vec![MediaItem::new("/tmp/baidu.png", "image")],
}],
)
.unwrap();
assert_eq!(message.id, "message-id");
assert_eq!(message.content, "这是已生成的截图");
assert_eq!(message.completion_status, CompletionStatus::Interrupted);
assert_eq!(message.media_refs.len(), 1);
}
}
use crate::bus::MessageBus;
use crate::providers::{LLMProvider, create_provider};
@ -710,10 +879,11 @@ impl Session {
self.id.to_string()
}
pub(super) fn add_message_in_memory(
pub(super) fn add_message_in_memory_with_version(
&mut self,
message: ChatMessage,
persist: bool,
advance_state_version: bool,
) -> Option<MessagePersistSnapshot> {
let is_user = message.role == "user";
let now = chrono::Utc::now().timestamp_millis();
@ -768,7 +938,9 @@ impl Session {
self.message_count += 1;
}
self.last_active_at = now;
self.state_version = self.state_version.wrapping_add(1);
if advance_state_version {
self.state_version = self.state_version.wrapping_add(1);
}
persist_snapshot.map(|(storage, session_id, msg_meta)| {
let session_meta = crate::storage::session::SessionMeta {
@ -797,7 +969,11 @@ impl Session {
/// Roll back messages that were appended in memory but whose atomic
/// persistence failed. This is only called while holding the session lock,
/// so the suffix check also protects against removing unrelated messages.
pub(super) fn rollback_message_suffix(&mut self, message_ids: &[String]) {
pub(super) fn rollback_message_suffix_with_version(
&mut self,
message_ids: &[String],
advance_state_version: bool,
) {
if message_ids.is_empty() || self.messages.len() < message_ids.len() {
return;
}
@ -819,7 +995,29 @@ impl Session {
self.seq_counter -= message_ids.len() as i64;
self.total_message_count -= message_ids.len() as i64;
self.message_count -= removed_user_messages;
self.state_version = self.state_version.wrapping_add(1);
if advance_state_version {
self.state_version = self.state_version.wrapping_add(1);
}
}
pub(super) fn owns_active_turn(&self, turn_id: &str) -> bool {
self.active_turn_emitter
.as_ref()
.is_some_and(|active| active.turn_id == turn_id)
}
#[cfg(test)]
pub(super) fn state_version_for_test(&self) -> u64 {
self.state_version
}
#[cfg(test)]
pub(super) fn set_active_turn_for_test(&mut self, turn_id: &str) {
let (_controller, emitter, _) = TurnController::start(self.id.to_string(), "test-message");
self.active_turn_emitter = Some(ActiveTurnEmitter {
turn_id: turn_id.to_string(),
emitter,
});
}
/// 获取消息历史
@ -2767,6 +2965,8 @@ fn spawn_agent_worker(
let commit_delivery = turn_delivery.clone();
let commit_target = turn_target.clone();
let turn_lifecycle = &turn_controller;
let pending_turn_deliveries = Arc::new(std::sync::Mutex::new(Vec::new()));
let scoped_turn_deliveries = pending_turn_deliveries.clone();
let process_future = async move {
let response_session_id = unified_str2.clone();
let process_result = crate::agent::sub_agent::DELEGATE_CONTEXT.scope(
@ -2777,7 +2977,7 @@ fn spawn_agent_worker(
},
agent.process_streaming(history_out.clone(), agent_turn.clone()),
).await;
let result = match process_result {
let mut result = match process_result {
Ok(r) => r,
Err(AgentError::LlmError(ref msg))
if is_context_overflow_error(msg) =>
@ -2921,6 +3121,8 @@ fn spawn_agent_worker(
}
};
let pending = take_current_turn_deliveries();
attach_pending_turn_deliveries(&mut result, pending);
let response_content = result.final_response.content;
let total_tokens = result.total_tokens;
let usage = result.usage;
@ -3005,15 +3207,20 @@ fn spawn_agent_worker(
let _ = bus2.publish_outbound(outbound).await;
}
};
let process_future =
CURRENT_TURN_ID.scope(Some(active_turn_id.clone()), process_future);
let process_future = CURRENT_TURN_DELIVERIES
.scope(Some(scoped_turn_deliveries), process_future);
tokio::select! {
() = process_future => {}
_ = cancel_rx => {
// cancelled — current_cancel already taken by /stop
let snapshot = turn_controller.snapshot();
if let Some(partial) = partial_assistant_message(
if let Some(partial) = partial_assistant_with_pending_deliveries(
&snapshot,
CompletionStatus::Cancelled,
take_pending_turn_deliveries(&pending_turn_deliveries),
) {
turn_controller.begin_finalizing();
match append_persisted_messages(&session, vec![partial]).await {

2
src/tools/mod.rs
View File

@ -41,7 +41,7 @@ pub use pty::{PtyManager, PtyTool};
pub use registry::ToolRegistry;
pub use send_message::SendMessageTool;
pub use todo::TodoTool;
pub use traits::{OutboundMessenger, Tool, ToolResult, ToolResultWithMedia};
pub use traits::{OutboundDelivery, OutboundMessenger, Tool, ToolResult, ToolResultWithMedia};
pub use web_fetch::WebFetchTool;
use crate::agent::SubAgentManager;

9
src/tools/send_message.rs
View File

@ -6,7 +6,7 @@ use mime_guess::mime;
use crate::bus::{MediaItem, MessageSource, SourceKind};
use super::traits::{OutboundMessenger, Tool, ToolResult};
use super::traits::{OutboundDelivery, OutboundMessenger, Tool, ToolResult};
pub struct SendMessageTool {
messenger: Arc<dyn OutboundMessenger>,
@ -144,11 +144,16 @@ target_chat_id 支持两种格式：<channel>:<chat_id>（发送到该聊天下
.send_message(channel, chat_id, dialog_id, content, source, media)
.await
{
Ok(()) => Ok(ToolResult {
Ok(OutboundDelivery::Delivered) => Ok(ToolResult {
success: true,
output: "消息已发送".to_string(),
error: None,
}),
Ok(OutboundDelivery::AttachedToCurrentTurn) => Ok(ToolResult {
success: true,
output: "附件已加入当前回复".to_string(),
error: None,
}),
Err(e) => Ok(ToolResult {
success: false,
output: String::new(),

8
src/tools/traits.rs
View File

@ -16,6 +16,12 @@ pub struct ToolResultWithMedia {
pub media_refs: Vec<MediaRef>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OutboundDelivery {
Delivered,
AttachedToCurrentTurn,
}
impl From<ToolResult> for ToolResultWithMedia {
fn from(result: ToolResult) -> Self {
Self {
@ -67,5 +73,5 @@ pub trait OutboundMessenger: Send + Sync {
content: &str,
source: MessageSource,
media: Vec<MediaItem>,
) -> Result<(), String>;
) -> Result<OutboundDelivery, String>;
}