From aa9f7a93db1a4f29de2020d6c0b0177619c7ef50 Mon Sep 17 00:00:00 2001 From: wsp1911 Date: Fri, 26 Jun 2026 11:35:13 +0800 Subject: [PATCH 1/4] fix(agentic): preserve remote session identity across cron and session tools - propagate remote workspace identity through agent turn submission and reply routing - resolve session-targeted cron and session tool actions from stored workspace bindings - keep scheduled-job reminders ordered after mode and remote file delivery reminders --- src/apps/desktop/src/api/agentic_api.rs | 6 + src/apps/desktop/src/api/miniapp_agent_api.rs | 2 + .../src/agentic/coordination/coordinator.rs | 228 ++++++++++++++-- .../src/agentic/coordination/scheduler.rs | 70 ++++- .../src/agentic/session/session_manager.rs | 243 +++++++++++++++--- .../tools/implementations/cron_tool.rs | 153 ++++++----- .../implementations/session_control_tool.rs | 107 +++++--- .../implementations/session_history_tool.rs | 10 +- .../implementations/session_message_tool.rs | 149 +++++++---- .../assembly/core/src/service/cron/service.rs | 188 +++++++++++++- .../core/src/service_agent_runtime.rs | 2 + src/crates/contracts/runtime-ports/src/lib.rs | 65 +++++ .../execution/agent-runtime/src/runtime.rs | 64 ++++- .../execution/agent-runtime/src/scheduler.rs | 4 + .../tests/scheduler_contracts.rs | 6 + .../interfaces/acp/src/runtime/prompt.rs | 4 + .../src/remote_connect.rs | 2 + 17 files changed, 1093 insertions(+), 210 deletions(-) diff --git a/src/apps/desktop/src/api/agentic_api.rs b/src/apps/desktop/src/api/agentic_api.rs index b0bacba51..e8d2b91e6 100644 --- a/src/apps/desktop/src/api/agentic_api.rs +++ b/src/apps/desktop/src/api/agentic_api.rs @@ -114,6 +114,8 @@ pub struct StartDialogTurnRequest { pub original_user_input: Option, pub agent_type: String, pub workspace_path: Option, + pub remote_connection_id: Option, + pub remote_ssh_host: Option, pub turn_id: Option, #[serde(default)] pub image_contexts: Option>, @@ -827,6 +829,8 @@ pub async fn start_dialog_turn( original_user_input, agent_type, workspace_path, + remote_connection_id, + remote_ssh_host, turn_id, image_contexts, user_message_metadata, @@ -851,6 +855,8 @@ pub async fn start_dialog_turn( turn_id, agent_type, workspace_path, + remote_connection_id, + remote_ssh_host, policy, None, user_message_metadata, diff --git a/src/apps/desktop/src/api/miniapp_agent_api.rs b/src/apps/desktop/src/api/miniapp_agent_api.rs index 662fd641b..4be06649a 100644 --- a/src/apps/desktop/src/api/miniapp_agent_api.rs +++ b/src/apps/desktop/src/api/miniapp_agent_api.rs @@ -258,6 +258,8 @@ pub async fn miniapp_agent_run( Some(submission_plan.run_id.clone()), MINIAPP_AGENT_KIND.to_string(), Some(submission_plan.workspace_path.clone()), + None, + None, policy, None, Some(submission_plan.metadata.clone()), diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index 2f028f844..5ca937a9e 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -30,6 +30,7 @@ use crate::agentic::goal_mode::{ }; use crate::agentic::image_analysis::ImageContextData; use crate::agentic::round_preempt::DialogRoundInjectionSource; +use crate::agentic::session::session_store_port::CoreSessionStorePort; use crate::agentic::session::SessionManager; use crate::agentic::side_question::build_btw_user_input; use crate::agentic::skill_agent_snapshot::{ @@ -59,9 +60,9 @@ use bitfun_agent_runtime::remote_file_delivery::{ TOOL_CONTEXT_REMOTE_FILE_DELIVERY_KEY, }; use bitfun_runtime_ports::{ - AgentBackgroundResultRequest, AgentThreadGoalDeliveryKind, AgentThreadGoalDeliveryRequest, - DelegationPolicy, SubagentContextMode, ThreadGoal, ThreadGoalContinuationPlan, - ThreadGoalStatus, + AgentBackgroundResultRequest, AgentSessionWorkspaceBinding, AgentThreadGoalDeliveryKind, + AgentThreadGoalDeliveryRequest, DelegationPolicy, SessionStoragePathRequest, SessionStorePort, + SubagentContextMode, ThreadGoal, ThreadGoalContinuationPlan, ThreadGoalStatus, }; use dashmap::DashMap; use log::{debug, error, info, warn}; @@ -515,7 +516,7 @@ pub struct ConversationCoordinator { } impl ConversationCoordinator { - async fn resolve_workspace_id_for_config(config: &SessionConfig) -> Option { + pub(crate) async fn resolve_workspace_id_for_config(config: &SessionConfig) -> Option { let explicit = config .workspace_id .as_deref() @@ -623,7 +624,9 @@ impl ConversationCoordinator { /// SSH connection (e.g. the user changed the port and the old ID is now /// stale), this method attempts to remap to the current workspace /// registration so that historical sessions continue to work. - async fn build_workspace_binding(config: &SessionConfig) -> Option { + pub(crate) async fn build_workspace_binding( + config: &SessionConfig, + ) -> Option { let workspace_path = config.workspace_path.as_ref()?; let path_buf = PathBuf::from(workspace_path); let workspace_id = Self::resolve_workspace_id_for_config(config).await; @@ -1896,6 +1899,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet Some(turn_id.clone()), ASSISTANT_BOOTSTRAP_AGENT_TYPE.to_string(), Some(workspace_root.to_string_lossy().to_string()), + None, + None, DialogSubmissionPolicy::for_source(DialogTriggerSource::DesktopApi) .with_skip_tool_confirmation(true), Some(metadata), @@ -1926,6 +1931,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet turn_id: Option, agent_type: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, submission_policy: DialogSubmissionPolicy, user_message_metadata: Option, ) -> BitFunResult<()> { @@ -1937,6 +1944,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet turn_id, agent_type, workspace_path, + remote_connection_id, + remote_ssh_host, submission_policy, user_message_metadata, Vec::new(), @@ -1954,6 +1963,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet turn_id: Option, agent_type: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, submission_policy: DialogSubmissionPolicy, user_message_metadata: Option, prepended_messages: Vec, @@ -1966,6 +1977,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet turn_id, agent_type, workspace_path, + remote_connection_id, + remote_ssh_host, submission_policy, user_message_metadata, prepended_messages, @@ -1984,6 +1997,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet turn_id: Option, agent_type: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, submission_policy: DialogSubmissionPolicy, user_message_metadata: Option, ) -> BitFunResult<()> { @@ -1995,6 +2010,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet turn_id, agent_type, workspace_path, + remote_connection_id, + remote_ssh_host, submission_policy, user_message_metadata, Vec::new(), @@ -2013,6 +2030,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet turn_id: Option, agent_type: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, submission_policy: DialogSubmissionPolicy, user_message_metadata: Option, prepended_messages: Vec, @@ -2025,6 +2044,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet turn_id, agent_type, workspace_path, + remote_connection_id, + remote_ssh_host, submission_policy, user_message_metadata, prepended_messages, @@ -2037,6 +2058,24 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet ThreadGoalStore::new(self.session_manager.as_ref()) } + async fn resolve_session_restore_path( + workspace_path: &str, + remote_connection_id: Option<&str>, + remote_ssh_host: Option<&str>, + ) -> BitFunResult { + let request = SessionStoragePathRequest { + workspace_path: PathBuf::from(workspace_path), + remote_connection_id: remote_connection_id.map(ToOwned::to_owned), + remote_ssh_host: remote_ssh_host.map(ToOwned::to_owned), + }; + + CoreSessionStorePort::default() + .resolve_session_storage_path(request) + .await + .map(|resolution| resolution.effective_storage_path) + .map_err(|error| BitFunError::Session(error.to_string())) + } + fn require_main_session_workspace(&self, session_id: &str) -> BitFunResult { let session = self .session_manager @@ -2716,6 +2755,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet turn_id: Option, agent_type: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, submission_policy: DialogSubmissionPolicy, extra_user_message_metadata: Option, additional_prepended_messages: Vec, @@ -2736,8 +2777,14 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet session_id )) })?; + let restore_path = Self::resolve_session_restore_path( + &workspace_path, + remote_connection_id.as_deref(), + remote_ssh_host.as_deref(), + ) + .await?; self.session_manager - .restore_session(Path::new(&workspace_path), &session_id) + .restore_session(&restore_path, &session_id) .await? } }; @@ -3004,14 +3051,11 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet ) .await?; let effective_user_input = wrapped_user_input_payload.content.clone(); - let mut prepended_messages = additional_prepended_messages; - if needs_computer_links_for_source(submission_policy.trigger_source) { - prepended_messages.push(Message::internal_reminder( - InternalReminderKind::RemoteFileDelivery, - remote_file_delivery_reminder(), - )); - } - prepended_messages.extend(wrapped_user_input_payload.prepended_messages.clone()); + let prepended_messages = merge_prepended_messages_for_turn( + additional_prepended_messages, + wrapped_user_input_payload.prepended_messages.clone(), + needs_computer_links_for_source(submission_policy.trigger_source), + ); if original_user_input != effective_user_input { let mut metadata = @@ -5138,6 +5182,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet Some(turn_id.clone()), child_session.agent_type.clone(), child_session.config.workspace_path.clone(), + child_session.config.remote_connection_id.clone(), + child_session.config.remote_ssh_host.clone(), DialogSubmissionPolicy::for_source(DialogTriggerSource::DesktopApi) .with_skip_tool_confirmation(true), user_message_metadata, @@ -5648,6 +5694,8 @@ impl bitfun_runtime_ports::AgentSubmissionPort for ConversationCoordinator { request.agent_type, SessionConfig { workspace_path: Some(workspace_path.clone()), + remote_connection_id: request.remote_connection_id.clone(), + remote_ssh_host: request.remote_ssh_host.clone(), ..Default::default() }, workspace_path, @@ -5705,6 +5753,8 @@ impl bitfun_runtime_ports::AgentSubmissionPort for ConversationCoordinator { Some(turn_id.clone()), session.agent_type.clone(), session.config.workspace_path.clone(), + session.config.remote_connection_id.clone(), + session.config.remote_ssh_host.clone(), DialogSubmissionPolicy::for_source(trigger_source), user_message_metadata, ) @@ -5726,10 +5776,27 @@ impl bitfun_runtime_ports::AgentSubmissionPort for ConversationCoordinator { &self, session_id: &str, ) -> bitfun_runtime_ports::PortResult> { - Ok(self + if let Some(session) = self.get_session_manager().get_session(session_id) { + return Ok(Some(session.agent_type.clone())); + } + + let Some(binding) = self .get_session_manager() - .get_session(session_id) - .map(|session| session.agent_type.clone())) + .resolve_session_workspace_binding(session_id) + .await + else { + return Ok(None); + }; + + self.restore_session(&binding.session_storage_path(), session_id) + .await + .map(|session| Some(session.agent_type)) + .map_err(|error| { + bitfun_runtime_ports::PortError::new( + bitfun_runtime_ports::PortErrorKind::Backend, + error.to_string(), + ) + }) } } @@ -5749,13 +5816,38 @@ fn runtime_session_summary(session: SessionSummary) -> bitfun_runtime_ports::Age } } +fn runtime_session_workspace_binding(binding: WorkspaceBinding) -> AgentSessionWorkspaceBinding { + AgentSessionWorkspaceBinding { + workspace_path: binding.root_path_string(), + remote_connection_id: binding.connection_id().map(ToOwned::to_owned), + remote_ssh_host: if binding.is_remote() { + Some(binding.session_identity.hostname.clone()).filter(|value| !value.trim().is_empty()) + } else { + None + }, + } +} + #[async_trait::async_trait] impl bitfun_runtime_ports::AgentSessionManagementPort for ConversationCoordinator { async fn list_sessions( &self, request: bitfun_runtime_ports::AgentSessionListRequest, ) -> bitfun_runtime_ports::PortResult> { - self.list_sessions(Path::new(&request.workspace_path)) + let effective_storage_path = Self::resolve_session_restore_path( + &request.workspace_path, + request.remote_connection_id.as_deref(), + request.remote_ssh_host.as_deref(), + ) + .await + .map_err(|error| { + bitfun_runtime_ports::PortError::new( + bitfun_runtime_ports::PortErrorKind::Backend, + error.to_string(), + ) + })?; + + self.list_sessions(&effective_storage_path) .await .map(|sessions| { sessions @@ -5775,7 +5867,20 @@ impl bitfun_runtime_ports::AgentSessionManagementPort for ConversationCoordinato &self, request: bitfun_runtime_ports::AgentSessionDeleteRequest, ) -> bitfun_runtime_ports::PortResult<()> { - self.delete_session(Path::new(&request.workspace_path), &request.session_id) + let effective_storage_path = Self::resolve_session_restore_path( + &request.workspace_path, + request.remote_connection_id.as_deref(), + request.remote_ssh_host.as_deref(), + ) + .await + .map_err(|error| { + bitfun_runtime_ports::PortError::new( + bitfun_runtime_ports::PortErrorKind::Backend, + error.to_string(), + ) + })?; + + self.delete_session(&effective_storage_path, &request.session_id) .await .map_err(|error| { bitfun_runtime_ports::PortError::new( @@ -5794,6 +5899,18 @@ impl bitfun_runtime_ports::AgentSessionManagementPort for ConversationCoordinato .await .map(|path| path.to_string_lossy().into_owned())) } + + async fn resolve_session_workspace_binding( + &self, + request: bitfun_runtime_ports::AgentSessionWorkspaceRequest, + ) -> bitfun_runtime_ports::PortResult> + { + Ok(self + .get_session_manager() + .resolve_session_workspace_binding(&request.session_id) + .await + .map(runtime_session_workspace_binding)) + } } #[async_trait::async_trait] @@ -5953,13 +6070,47 @@ pub fn get_global_coordinator() -> Option> { GLOBAL_COORDINATOR.get().cloned() } +fn merge_prepended_messages_for_turn( + additional_prepended_messages: Vec, + wrapped_prepended_messages: Vec, + include_remote_file_delivery: bool, +) -> Vec { + let mut prepended_messages = Vec::new(); + let mut scheduled_job_messages = Vec::new(); + let mut remote_file_delivery_messages = Vec::new(); + + for message in additional_prepended_messages { + if matches!( + message.internal_reminder_kind(), + Some(InternalReminderKind::ScheduledJob) + ) { + scheduled_job_messages.push(message); + } else { + prepended_messages.push(message); + } + } + + if include_remote_file_delivery { + remote_file_delivery_messages.push(Message::internal_reminder( + InternalReminderKind::RemoteFileDelivery, + remote_file_delivery_reminder(), + )); + } + + prepended_messages.extend(wrapped_prepended_messages); + prepended_messages.extend(remote_file_delivery_messages); + prepended_messages.extend(scheduled_job_messages); + prepended_messages +} + #[cfg(test)] mod tests { use super::{ - normalize_subagent_max_concurrency, resolve_agent_session_create_created_by, - resolve_agent_submission_turn_id, ConversationCoordinator, + merge_prepended_messages_for_turn, normalize_subagent_max_concurrency, + resolve_agent_session_create_created_by, resolve_agent_submission_turn_id, + ConversationCoordinator, }; - use crate::agentic::core::SessionConfig; + use crate::agentic::core::{InternalReminderKind, Message, SessionConfig}; use crate::agentic::events::{EventQueue, EventQueueConfig, EventRouter}; use crate::agentic::execution::{ ExecutionEngine, ExecutionEngineConfig, RoundExecutor, StreamProcessor, @@ -6214,6 +6365,8 @@ mod tests { session_name: "Worker".to_string(), agent_type: "agentic".to_string(), workspace_path: Some(workspace_path.to_string_lossy().into_owned()), + remote_connection_id: None, + remote_ssh_host: None, metadata, }, ) @@ -6356,4 +6509,35 @@ mod tests { Some(baseline_snapshot) ); } + + #[test] + fn merge_prepended_messages_places_scheduled_job_after_mode_reminder() { + let merged = merge_prepended_messages_for_turn( + vec![ + Message::internal_reminder(InternalReminderKind::ScheduledJob, "scheduled"), + Message::internal_reminder(InternalReminderKind::Generic, "generic"), + ], + vec![ + Message::internal_reminder(InternalReminderKind::SkillListingDiff, "skills"), + Message::internal_reminder(InternalReminderKind::AgentMode, "mode"), + ], + true, + ); + + let kinds = merged + .iter() + .map(|message| message.internal_reminder_kind()) + .collect::>(); + + assert_eq!( + kinds, + vec![ + Some(InternalReminderKind::Generic), + Some(InternalReminderKind::SkillListingDiff), + Some(InternalReminderKind::AgentMode), + Some(InternalReminderKind::RemoteFileDelivery), + Some(InternalReminderKind::ScheduledJob), + ] + ); + } } diff --git a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs index 3a7ef5b2a..1517eaf04 100644 --- a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs +++ b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs @@ -20,10 +20,11 @@ use crate::agentic::goal_mode::{ use crate::agentic::image_analysis::ImageContextData; use crate::agentic::init_agents_md::build_init_agents_md_user_input; use crate::agentic::round_preempt::{DialogRoundInjectionSource, SessionRoundInjectionBuffer}; +use crate::agentic::session::session_store_port::CoreSessionStorePort; use crate::agentic::session::SessionManager; use bitfun_runtime_ports::{ThreadGoal, MAX_THREAD_GOAL_AUTO_CONTINUATIONS}; use log::{debug, info, warn}; -use std::path::Path; +use std::path::PathBuf; use std::sync::Arc; use std::sync::OnceLock; use std::time::{Duration, SystemTime}; @@ -48,6 +49,7 @@ use bitfun_runtime_ports::{ AgentThreadGoalDeliveryKind, AgentThreadGoalDeliveryRequest, AgentTurnCancellationPort, AgentTurnCancellationRequest, AgentTurnCancellationResult, DialogSessionStateFact, DialogSubmitQueueAction, DialogSubmitQueueFacts, PortError, PortErrorKind, PortResult, + SessionStoragePathRequest, SessionStorePort, }; pub use bitfun_runtime_ports::{ AgentSessionReplyRoute, DialogQueuePriority, DialogSteerOutcome, DialogSubmissionPolicy, @@ -63,6 +65,8 @@ pub struct QueuedTurn { pub turn_id: Option, pub agent_type: String, pub workspace_path: Option, + pub remote_connection_id: Option, + pub remote_ssh_host: Option, pub policy: DialogSubmissionPolicy, pub reply_route: Option, pub user_message_metadata: Option, @@ -237,6 +241,8 @@ impl DialogScheduler { None, agent_type, workspace_path, + None, + None, DialogSubmissionPolicy::new( DialogTriggerSource::AgentSession, queue_priority, @@ -295,6 +301,8 @@ impl DialogScheduler { None, agent_type, workspace_path, + None, + None, DialogSubmissionPolicy::new( DialogTriggerSource::AgentSession, queue_priority, @@ -358,6 +366,8 @@ impl DialogScheduler { None, agent_type, workspace_path, + None, + None, DialogSubmissionPolicy::new( DialogTriggerSource::AgentSession, queue_priority, @@ -379,7 +389,7 @@ impl DialogScheduler { policy: DialogSubmissionPolicy, ) -> Result { let agent_type = self - .resolve_session_agent_type(&session_id, workspace_path.as_deref()) + .resolve_session_agent_type(&session_id, workspace_path.as_deref(), None, None) .await?; let (user_input, prepended_messages) = build_init_agents_md_user_input() .await @@ -392,6 +402,8 @@ impl DialogScheduler { None, agent_type, workspace_path, + None, + None, policy, None, None, @@ -428,6 +440,8 @@ impl DialogScheduler { turn_id: Option, agent_type: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, policy: DialogSubmissionPolicy, reply_route: Option, user_message_metadata: Option, @@ -440,6 +454,8 @@ impl DialogScheduler { turn_id, agent_type, workspace_path, + remote_connection_id, + remote_ssh_host, policy, reply_route, user_message_metadata, @@ -458,6 +474,8 @@ impl DialogScheduler { turn_id: Option, agent_type: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, policy: DialogSubmissionPolicy, reply_route: Option, user_message_metadata: Option, @@ -472,6 +490,8 @@ impl DialogScheduler { turn_id: Some(resolved_turn_id.clone()), agent_type, workspace_path, + remote_connection_id, + remote_ssh_host, policy, reply_route, user_message_metadata, @@ -486,6 +506,8 @@ impl DialogScheduler { &self, session_id: &str, workspace_path: Option<&str>, + remote_connection_id: Option<&str>, + remote_ssh_host: Option<&str>, ) -> Result { let session = match self.session_manager.get_session(session_id) { Some(session) => session, @@ -496,8 +518,14 @@ impl DialogScheduler { session_id ) })?; + let restore_path = Self::resolve_session_restore_path( + workspace_path, + remote_connection_id, + remote_ssh_host, + ) + .await?; self.session_manager - .restore_session(Path::new(workspace_path), session_id) + .restore_session(&restore_path, session_id) .await .map_err(|error| error.to_string())? } @@ -510,6 +538,24 @@ impl DialogScheduler { } } + async fn resolve_session_restore_path( + workspace_path: &str, + remote_connection_id: Option<&str>, + remote_ssh_host: Option<&str>, + ) -> Result { + let request = SessionStoragePathRequest { + workspace_path: PathBuf::from(workspace_path), + remote_connection_id: remote_connection_id.map(ToOwned::to_owned), + remote_ssh_host: remote_ssh_host.map(ToOwned::to_owned), + }; + + CoreSessionStorePort::default() + .resolve_session_storage_path(request) + .await + .map(|resolution| resolution.effective_storage_path) + .map_err(|error| error.to_string()) + } + async fn submit_queued_turn( &self, session_id: String, @@ -742,6 +788,8 @@ impl DialogScheduler { queued_turn.turn_id.clone(), queued_turn.agent_type.clone(), queued_turn.workspace_path.clone(), + queued_turn.remote_connection_id.clone(), + queued_turn.remote_ssh_host.clone(), queued_turn.policy, queued_turn.user_message_metadata.clone(), ) @@ -756,6 +804,8 @@ impl DialogScheduler { queued_turn.turn_id.clone(), queued_turn.agent_type.clone(), queued_turn.workspace_path.clone(), + queued_turn.remote_connection_id.clone(), + queued_turn.remote_ssh_host.clone(), queued_turn.policy, queued_turn.user_message_metadata.clone(), queued_turn.prepended_messages.clone(), @@ -774,6 +824,8 @@ impl DialogScheduler { queued_turn.turn_id.clone(), queued_turn.agent_type.clone(), queued_turn.workspace_path.clone(), + queued_turn.remote_connection_id.clone(), + queued_turn.remote_ssh_host.clone(), queued_turn.policy, queued_turn.user_message_metadata.clone(), ) @@ -791,6 +843,8 @@ impl DialogScheduler { queued_turn.turn_id.clone(), queued_turn.agent_type.clone(), queued_turn.workspace_path.clone(), + queued_turn.remote_connection_id.clone(), + queued_turn.remote_ssh_host.clone(), queued_turn.policy, queued_turn.user_message_metadata.clone(), queued_turn.prepended_messages.clone(), @@ -844,6 +898,8 @@ impl DialogScheduler { let reply_user_input = plan.user_input; let target_session_id = plan.target_session_id; let target_workspace_path = plan.target_workspace_path; + let target_remote_connection_id = plan.target_remote_connection_id; + let target_remote_ssh_host = plan.target_remote_ssh_host; let prepended_messages = vec![Message::internal_reminder( InternalReminderKind::SessionMessageReply, plan.reminder_text, @@ -857,6 +913,8 @@ impl DialogScheduler { None, String::new(), Some(target_workspace_path), + target_remote_connection_id, + target_remote_ssh_host, DialogSubmissionPolicy::for_source(DialogTriggerSource::AgentSession), None, None, @@ -985,6 +1043,8 @@ impl DialogScheduler { None, active_turn.agent_type_owned(), active_turn.workspace_path_owned(), + None, + None, DialogSubmissionPolicy::for_source( DialogTriggerSource::AgentSession, ), @@ -1205,6 +1265,8 @@ impl AgentDialogTurnPort for DialogScheduler { request.turn_id, request.agent_type, request.workspace_path, + request.remote_connection_id, + request.remote_ssh_host, request.policy, request.reply_route, user_message_metadata, @@ -1363,6 +1425,8 @@ mod tests { Some(AgentSessionReplyRoute { source_session_id: source_session_id.to_string(), source_workspace_path: "/source".to_string(), + source_remote_connection_id: None, + source_remote_ssh_host: None, }), ) } diff --git a/src/crates/assembly/core/src/agentic/session/session_manager.rs b/src/crates/assembly/core/src/agentic/session/session_manager.rs index ad5290425..a6e973324 100644 --- a/src/crates/assembly/core/src/agentic/session/session_manager.rs +++ b/src/crates/assembly/core/src/agentic/session/session_manager.rs @@ -20,17 +20,20 @@ use crate::agentic::session::{ TurnSkillAgentSnapshotStore, UserContextCacheIdentity, }; use crate::agentic::skill_agent_snapshot::TurnSkillAgentSnapshot; +use crate::agentic::workspace::WorkspaceBinding; +use crate::agentic::ConversationCoordinator; use crate::infrastructure::ai::get_global_ai_client_factory; use crate::service::config::{ get_app_language_code, get_global_config_service, short_model_user_language_instruction, subscribe_config_updates, ConfigUpdateEvent, }; +use crate::service::remote_ssh::workspace_state::LOCAL_WORKSPACE_SSH_HOST; use crate::service::session::{ DialogTurnData, DialogTurnKind, ModelRoundData, SessionMetadata, SessionRelationship, TextItemData, TurnStatus, UserMessageData, }; use crate::service::snapshot::ensure_snapshot_manager_for_workspace; -use crate::service::workspace::get_global_workspace_service; +use crate::service::workspace::{get_global_workspace_service, WorkspaceInfo, WorkspaceKind}; use crate::util::errors::{BitFunError, BitFunResult}; use crate::util::sanitize_plain_model_output; use crate::util::timing::elapsed_ms_u64; @@ -107,7 +110,7 @@ pub struct SessionManager { /// Active sessions in memory sessions: Arc>, - /// Runtime cache of session_id -> effective workspace path. + /// Runtime cache of session_id -> effective session storage path. /// Populated on session create/restore and used to restore evicted sessions /// or resolve workspace-bound operations that only receive a session_id. /// This cache is intentionally retained across memory eviction, but should @@ -477,43 +480,223 @@ impl SessionManager { } } - let workspace_service = get_global_workspace_service()?; - let mut workspaces = workspace_service.list_workspace_infos().await; - workspaces.sort_by(|left, right| right.last_accessed.cmp(&left.last_accessed)); - let candidates: Vec = workspaces - .into_iter() - .map(|workspace| workspace.root_path) - .collect(); + if let Some(binding) = self.resolve_session_workspace_binding(session_id).await { + return Some(binding.root_path().to_path_buf()); + } - for workspace_path in candidates { - match self - .persistence_manager - .load_session_metadata(&workspace_path, session_id) + None + } + + pub async fn resolve_session_workspace_binding( + &self, + session_id: &str, + ) -> Option { + if let Some(config) = self + .get_session(session_id) + .map(|session| session.config.clone()) + { + if let Some(binding) = ConversationCoordinator::build_workspace_binding(&config).await { + return Some(binding); + } + } + + let indexed_workspace_path = self + .session_workspace_index + .get(session_id) + .map(|entry| entry.clone()); + if let Some(session_storage_path) = indexed_workspace_path { + if let Some(binding) = self + .resolve_persisted_session_workspace_binding( + session_id, + &session_storage_path, + None, + ) .await { - Ok(Some(metadata)) => { - if let Some(bound_workspace) = - metadata.workspace_path.filter(|path| !path.is_empty()) - { - return Some(PathBuf::from(bound_workspace)); - } - return Some(workspace_path); - } - Ok(None) => {} - Err(err) => { - debug!( - "Failed to load session metadata while resolving workspace: session_id={} workspace={} error={}", - session_id, - workspace_path.display(), - err - ); - } + return Some(binding); + } + } + + for workspace in self.tracked_workspace_candidates().await? { + let Some(session_storage_path) = + Self::session_storage_path_for_workspace_info(&workspace).await + else { + continue; + }; + + if let Some(binding) = self + .resolve_persisted_session_workspace_binding( + session_id, + &session_storage_path, + Some(&workspace), + ) + .await + { + self.session_workspace_index + .insert(session_id.to_string(), session_storage_path); + return Some(binding); } } None } + async fn resolve_persisted_session_workspace_binding( + &self, + session_id: &str, + session_storage_path: &Path, + workspace_hint: Option<&WorkspaceInfo>, + ) -> Option { + let metadata = match self + .persistence_manager + .load_session_metadata(session_storage_path, session_id) + .await + { + Ok(Some(metadata)) => metadata, + Ok(None) => return None, + Err(err) => { + debug!( + "Failed to load session metadata while resolving workspace binding: session_id={} storage_path={} error={}", + session_id, + session_storage_path.display(), + err + ); + return None; + } + }; + + let config = self + .session_config_from_persisted_metadata(&metadata, workspace_hint) + .await?; + + ConversationCoordinator::build_workspace_binding(&config).await + } + + async fn session_config_from_persisted_metadata( + &self, + metadata: &SessionMetadata, + workspace_hint: Option<&WorkspaceInfo>, + ) -> Option { + let workspace_path = metadata + .workspace_path + .as_deref() + .map(str::trim) + .filter(|path| !path.is_empty()) + .map(str::to_string) + .or_else(|| { + workspace_hint.map(|workspace| workspace.root_path.to_string_lossy().to_string()) + })?; + + let mut config = SessionConfig { + workspace_path: Some(workspace_path.clone()), + ..SessionConfig::default() + }; + + let remote_hostname = metadata + .workspace_hostname + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty() && *value != LOCAL_WORKSPACE_SSH_HOST) + .map(str::to_string); + + let matched_workspace = match workspace_hint { + Some(workspace) => Some(workspace.clone()), + None if remote_hostname.is_some() => { + self.match_tracked_remote_workspace(&workspace_path, remote_hostname.as_deref()) + .await + } + None => None, + }; + + if let Some(workspace) = matched_workspace.as_ref() { + config.workspace_id = Some(workspace.id.clone()); + if workspace.workspace_kind == WorkspaceKind::Remote { + config.remote_connection_id = + workspace.remote_ssh_connection_id().map(ToOwned::to_owned); + config.remote_ssh_host = workspace + .metadata + .get("sshHost") + .and_then(|value| value.as_str()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string); + if config.remote_connection_id.is_none() { + return None; + } + } + } else if remote_hostname.is_some() { + return None; + } + + Some(config) + } + + async fn match_tracked_remote_workspace( + &self, + workspace_path: &str, + ssh_host: Option<&str>, + ) -> Option { + let Some(ssh_host) = ssh_host.map(str::trim).filter(|value| !value.is_empty()) else { + return None; + }; + + let normalized_workspace_path = + crate::service::remote_ssh::normalize_remote_workspace_path(workspace_path); + + self.tracked_workspace_candidates() + .await? + .into_iter() + .find(|workspace| { + if workspace.workspace_kind != WorkspaceKind::Remote { + return false; + } + + if crate::service::remote_ssh::normalize_remote_workspace_path( + &workspace.root_path.to_string_lossy(), + ) != normalized_workspace_path + { + return false; + } + + let workspace_host = workspace + .metadata + .get("sshHost") + .and_then(|value| value.as_str()) + .map(str::trim) + .filter(|value| !value.is_empty()); + + workspace_host == Some(ssh_host) + }) + } + + async fn tracked_workspace_candidates(&self) -> Option> { + let workspace_service = get_global_workspace_service()?; + let mut workspaces = workspace_service.list_workspace_infos().await; + workspaces.sort_by(|left, right| right.last_accessed.cmp(&left.last_accessed)); + Some(workspaces) + } + + async fn session_storage_path_for_workspace_info(workspace: &WorkspaceInfo) -> Option { + let remote_connection_id = workspace.remote_ssh_connection_id().map(ToOwned::to_owned); + let remote_ssh_host = workspace + .metadata + .get("sshHost") + .and_then(|value| value.as_str()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string); + + CoreSessionStorePort::default() + .resolve_session_storage_path(SessionStoragePathRequest { + workspace_path: workspace.root_path.clone(), + remote_connection_id, + remote_ssh_host, + }) + .await + .ok() + .map(|resolution| resolution.effective_storage_path) + } + fn build_messages_from_turns(turns: &[DialogTurnData]) -> Vec { let mut messages = Vec::new(); diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/cron_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/cron_tool.rs index 7f2fc7682..56bbd28c1 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/cron_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/cron_tool.rs @@ -4,6 +4,7 @@ use crate::agentic::tools::framework::{ Tool, ToolExposure, ToolRenderOptions, ToolResult, ToolUseContext, ValidationResult, }; use crate::agentic::tools::workspace_paths::posix_style_path_is_absolute; +use crate::agentic::workspace::WorkspaceBinding; use crate::service::{ cron::{ CreateCronJobRequest, CronJob, CronJobPayload, CronJobRunStatus, CronJobTarget, @@ -11,15 +12,12 @@ use crate::service::{ }, get_global_cron_service, }; -use crate::service_agent_runtime::CoreServiceAgentRuntime; use crate::util::errors::{BitFunError, BitFunResult}; use async_trait::async_trait; -use bitfun_agent_runtime::sdk::AgentRuntime; -use bitfun_runtime_ports::{AgentSessionListRequest, AgentSessionWorkspaceRequest}; use chrono::{DateTime, Local, SecondsFormat, TimeZone}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; -use std::path::Path; +use std::path::{Path, PathBuf}; const DEFAULT_JOB_NAME: &str = "Cron job"; @@ -120,41 +118,27 @@ impl CronTool { self.resolve_workspace(workspace.to_string_lossy().as_ref(), Some(context)) } - fn agent_runtime_if_available(&self) -> BitFunResult> { - let Some(coordinator) = get_global_coordinator() else { - return Ok(None); - }; - CoreServiceAgentRuntime::agent_runtime(coordinator) - .map(Some) - .map_err(BitFunError::tool) - } - - fn require_agent_runtime(&self) -> BitFunResult { - self.agent_runtime_if_available()? - .ok_or_else(|| BitFunError::tool("coordinator not initialized".to_string())) - } - async fn resolve_effective_workspace_for_session( &self, session_id: &str, context: &ToolUseContext, - ) -> BitFunResult { - if let Some(runtime) = self.agent_runtime_if_available()? { - if let Some(resolved) = runtime - .resolve_session_workspace_path(AgentSessionWorkspaceRequest { - session_id: session_id.to_string(), - }) + ) -> BitFunResult { + if let Some(coordinator) = get_global_coordinator() { + if let Some(resolved) = coordinator + .get_session_manager() + .resolve_session_workspace_binding(session_id) .await - .map_err(|error| { - BitFunError::tool(CoreServiceAgentRuntime::runtime_error_message(error)) - })? { return Ok(resolved); } } if context.session_id.as_deref() == Some(session_id) { - return self.resolve_workspace_from_context(context); + if let Some(binding) = context.workspace.as_ref() { + return Ok(binding.clone()); + } + let resolved = self.resolve_workspace_from_context(context)?; + return Ok(WorkspaceBinding::new(None, PathBuf::from(resolved))); } Err(BitFunError::tool(format!( @@ -182,29 +166,6 @@ impl CronTool { Ok(resolved) } - async fn ensure_session_exists(&self, workspace: &str, session_id: &str) -> BitFunResult<()> { - let runtime = self.require_agent_runtime()?; - let sessions = runtime - .list_sessions(AgentSessionListRequest { - workspace_path: workspace.to_string(), - }) - .await - .map_err(|error| { - BitFunError::tool(CoreServiceAgentRuntime::runtime_error_message(error)) - })?; - if sessions - .iter() - .any(|session| session.session_id == session_id) - { - return Ok(()); - } - - Err(BitFunError::NotFound(format!( - "Session '{}' not found in workspace '{}'", - session_id, workspace - ))) - } - fn normalize_add_name(name: Option) -> String { match name { Some(name) if !name.trim().is_empty() => name.trim().to_string(), @@ -212,6 +173,20 @@ impl CronTool { } } + fn workspace_ref_from_binding(&self, binding: &WorkspaceBinding) -> CronWorkspaceRef { + CronWorkspaceRef { + workspace_id: binding.workspace_id.clone(), + workspace_path: binding.root_path_string(), + remote_connection_id: binding.connection_id().map(ToOwned::to_owned), + remote_ssh_host: if binding.is_remote() { + Some(binding.session_identity.hostname.clone()) + .filter(|value| !value.trim().is_empty()) + } else { + None + }, + } + } + fn normalize_optional_name(name: Option) -> BitFunResult> { match name { Some(name) if name.trim().is_empty() => Err(BitFunError::tool( @@ -998,14 +973,16 @@ Patch schema for "update": .ok_or_else(|| BitFunError::tool("cron service not initialized".to_string()))?; let session_id = self.resolve_effective_session_id(params.session_id.as_deref(), context)?; - let workspace = self + let workspace_binding = self .resolve_effective_workspace_for_session(&session_id, context) .await?; + let workspace = workspace_binding.root_path_string(); + let workspace_ref = self.workspace_ref_from_binding(&workspace_binding); let mut jobs = cron_service .list_jobs_filtered( - Some(&workspace), - None, - None, + Some(&workspace_ref.workspace_path), + workspace_ref.workspace_id.as_deref(), + workspace_ref.remote_connection_id.as_deref(), Some(&session_id), Some(CronJobTargetKind::Session), ) @@ -1038,15 +1015,16 @@ Patch schema for "update": .ok_or_else(|| BitFunError::tool("cron service not initialized".to_string()))?; let session_id = self.resolve_effective_session_id(params.session_id.as_deref(), context)?; - let workspace = self + let workspace_binding = self .resolve_effective_workspace_for_session(&session_id, context) .await?; + let workspace = workspace_binding.root_path_string(); let job = params .job .ok_or_else(|| BitFunError::tool("job is required for add".to_string()))?; Self::validate_payload(&job.payload, "job")?; - self.ensure_session_exists(&workspace, &session_id).await?; + let target_workspace = self.workspace_ref_from_binding(&workspace_binding); let created = cron_service .create_job(CreateCronJobRequest { @@ -1056,12 +1034,7 @@ Patch schema for "update": enabled: job.enabled.unwrap_or(true), target: CronJobTarget::Session { session_id: session_id.clone(), - workspace: CronWorkspaceRef { - workspace_id: None, - workspace_path: workspace.clone(), - remote_connection_id: None, - remote_ssh_host: None, - }, + workspace: target_workspace, }, }) .await?; @@ -1196,8 +1169,11 @@ Patch schema for "update": mod tests { use super::*; use crate::agentic::tools::framework::ToolUseContext; + use crate::agentic::workspace::WorkspaceBinding; + use crate::service::remote_ssh::workspace_state::workspace_session_identity; use serde_json::json; use std::collections::HashMap; + use std::path::PathBuf; fn empty_context() -> ToolUseContext { ToolUseContext { @@ -1214,6 +1190,33 @@ mod tests { } } + fn remote_context( + root: &str, + workspace_id: Option, + session_id: Option<&str>, + ) -> ToolUseContext { + let session_identity = workspace_session_identity(root, Some("conn-1"), Some("ssh.dev")) + .expect("remote identity"); + ToolUseContext { + tool_call_id: None, + agent_type: None, + session_id: session_id.map(ToOwned::to_owned), + dialog_turn_id: None, + workspace: Some(WorkspaceBinding::new_remote( + workspace_id, + PathBuf::from(root), + "conn-1".to_string(), + "Dev SSH".to_string(), + session_identity, + )), + unlocked_collapsed_tools: Vec::new(), + custom_data: HashMap::new(), + computer_use_host: None, + runtime_tool_restrictions: Default::default(), + runtime_handles: bitfun_runtime_ports::ToolRuntimeHandles::default(), + } + } + #[tokio::test] async fn validate_list_allows_missing_workspace_when_session_id_present() { let tool = CronTool::new(); @@ -1277,4 +1280,28 @@ mod tests { .unwrap_or_default() .contains("unknown field")); } + + #[test] + fn workspace_ref_from_binding_uses_remote_context_identity() { + let tool = CronTool::new(); + let context = remote_context( + "/home/wsp/projects/test", + Some("remote_workspace_1".to_string()), + Some("session-1"), + ); + + let workspace_ref = + tool.workspace_ref_from_binding(context.workspace.as_ref().expect("workspace binding")); + + assert_eq!( + workspace_ref.workspace_id.as_deref(), + Some("remote_workspace_1") + ); + assert_eq!(workspace_ref.workspace_path, "/home/wsp/projects/test"); + assert_eq!( + workspace_ref.remote_connection_id.as_deref(), + Some("conn-1") + ); + assert_eq!(workspace_ref.remote_ssh_host.as_deref(), Some("ssh.dev")); + } } diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/session_control_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/session_control_tool.rs index ead63e15e..129857984 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/session_control_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/session_control_tool.rs @@ -24,8 +24,8 @@ use bitfun_agent_runtime::session_control::{ }; use bitfun_runtime_ports::{ AgentSessionCreateRequest, AgentSessionDeleteRequest, AgentSessionListRequest, - AgentSessionSummary, AgentSessionWorkspaceRequest, AgentSubmissionSource, - AgentTurnCancellationRequest, + AgentSessionSummary, AgentSessionWorkspaceBinding, AgentSessionWorkspaceRequest, + AgentSubmissionSource, AgentTurnCancellationRequest, }; use serde_json::{json, Value}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; @@ -35,6 +35,13 @@ pub struct SessionControlTool; const CANCEL_WAIT_TIMEOUT: Duration = Duration::from_secs(3); +#[derive(Debug, Clone)] +struct SessionControlWorkspaceTarget { + display_workspace: String, + remote_connection_id: Option, + remote_ssh_host: Option, +} + impl Default for SessionControlTool { fn default() -> Self { Self::new() @@ -88,14 +95,14 @@ impl SessionControlTool { session_id: Option<&str>, context: &ToolUseContext, runtime: &AgentRuntime, - ) -> BitFunResult { + ) -> BitFunResult { match action { SessionControlAction::Cancel | SessionControlAction::Delete => { let session_id = session_id.ok_or_else(|| { BitFunError::tool(format!("session_id is required for {}", action.as_str())) })?; - if let Some(resolved) = runtime - .resolve_session_workspace_path(AgentSessionWorkspaceRequest { + if let Some(binding) = runtime + .resolve_session_workspace_binding(AgentSessionWorkspaceRequest { session_id: session_id.to_string(), }) .await @@ -103,22 +110,47 @@ impl SessionControlTool { BitFunError::tool(CoreServiceAgentRuntime::runtime_error_message(error)) })? { - return Ok(resolved); + return Ok(Self::workspace_target_from_binding(binding)); } Err(BitFunError::NotFound(format!( "Workspace for session '{}' could not be resolved", session_id ))) } - SessionControlAction::Create | SessionControlAction::List => context - .workspace_root() - .map(|path| normalize_path(path.to_string_lossy().as_ref())) - .ok_or_else(|| { + SessionControlAction::Create | SessionControlAction::List => { + let workspace = context.workspace.as_ref().ok_or_else(|| { BitFunError::tool(format!( "workspace is required for {} when the current workspace is unavailable", action.as_str() )) - }), + })?; + Ok(Self::workspace_target_from_context(workspace)) + } + } + } + + fn workspace_target_from_context( + workspace: &crate::agentic::WorkspaceBinding, + ) -> SessionControlWorkspaceTarget { + SessionControlWorkspaceTarget { + display_workspace: normalize_path(&workspace.root_path_string()), + remote_connection_id: workspace.connection_id().map(ToOwned::to_owned), + remote_ssh_host: if workspace.is_remote() { + Some(workspace.session_identity.hostname.clone()) + .filter(|value| !value.trim().is_empty()) + } else { + None + }, + } + } + + fn workspace_target_from_binding( + binding: AgentSessionWorkspaceBinding, + ) -> SessionControlWorkspaceTarget { + SessionControlWorkspaceTarget { + display_workspace: binding.workspace_path, + remote_connection_id: binding.remote_connection_id, + remote_ssh_host: binding.remote_ssh_host, } } @@ -141,12 +173,14 @@ impl SessionControlTool { async fn ensure_session_exists( &self, runtime: &AgentRuntime, - workspace: &str, + workspace: &SessionControlWorkspaceTarget, session_id: &str, ) -> BitFunResult<()> { let existing_sessions = runtime .list_sessions(AgentSessionListRequest { - workspace_path: workspace.to_string(), + workspace_path: workspace.display_workspace.clone(), + remote_connection_id: workspace.remote_connection_id.clone(), + remote_ssh_host: workspace.remote_ssh_host.clone(), }) .await .map_err(|error| { @@ -160,7 +194,7 @@ impl SessionControlTool { } else { Err(BitFunError::NotFound(format!( "Session '{}' not found in workspace '{}'", - session_id, workspace + session_id, workspace.display_workspace ))) } } @@ -344,7 +378,9 @@ Arguments: .create_session(AgentSessionCreateRequest { session_name, agent_type, - workspace_path: Some(workspace.clone()), + workspace_path: Some(workspace.display_workspace.clone()), + remote_connection_id: workspace.remote_connection_id.clone(), + remote_ssh_host: workspace.remote_ssh_host.clone(), metadata, }) .await @@ -356,7 +392,7 @@ Arguments: let created_agent_type = session.agent_type.clone(); let result_for_assistant = session_control_created_result_message( &created_session_id, - &workspace, + &workspace.display_workspace, &created_agent_type, ); @@ -364,7 +400,7 @@ Arguments: data: json!({ "success": true, "action": "create", - "workspace": workspace.clone(), + "workspace": workspace.display_workspace.clone(), "session": { "session_id": created_session_id, "session_name": created_session_name, @@ -388,7 +424,9 @@ Arguments: &runtime, ) .await?; - if self.current_workspace_session(context, &workspace) == Some(session_id) { + if self.current_workspace_session(context, &workspace.display_workspace) + == Some(session_id) + { return Err(BitFunError::tool( "cannot cancel the current session from SessionControl".to_string(), )); @@ -440,7 +478,7 @@ Arguments: let status = session_control_cancel_status(cancelled_turn_id.as_deref()); let result_for_assistant = session_control_cancel_result_message( session_id, - &workspace, + &workspace.display_workspace, cancelled_turn_id.as_deref(), ); @@ -448,7 +486,7 @@ Arguments: data: json!({ "success": true, "action": "cancel", - "workspace": workspace.clone(), + "workspace": workspace.display_workspace.clone(), "session_id": session_id, "had_active_turn": had_active_turn, "cancelled_turn_id": cancelled_turn_id, @@ -471,7 +509,9 @@ Arguments: &runtime, ) .await?; - if self.current_workspace_session(context, &workspace) == Some(session_id) { + if self.current_workspace_session(context, &workspace.display_workspace) + == Some(session_id) + { return Err(BitFunError::tool( "cannot delete the current session from SessionControl".to_string(), )); @@ -482,8 +522,10 @@ Arguments: runtime .delete_session(AgentSessionDeleteRequest { - workspace_path: workspace.clone(), + workspace_path: workspace.display_workspace.clone(), session_id: session_id.to_string(), + remote_connection_id: workspace.remote_connection_id.clone(), + remote_ssh_host: workspace.remote_ssh_host.clone(), }) .await .map_err(|error| { @@ -494,11 +536,12 @@ Arguments: data: json!({ "success": true, "action": "delete", - "workspace": workspace.clone(), + "workspace": workspace.display_workspace.clone(), "session_id": session_id, }), result_for_assistant: Some(session_control_deleted_result_message( - session_id, &workspace, + session_id, + &workspace.display_workspace, )), image_attachments: None, }]) @@ -514,21 +557,27 @@ Arguments: .await?; let sessions = runtime .list_sessions(AgentSessionListRequest { - workspace_path: workspace.clone(), + workspace_path: workspace.display_workspace.clone(), + remote_connection_id: workspace.remote_connection_id.clone(), + remote_ssh_host: workspace.remote_ssh_host.clone(), }) .await .map_err(|error| { BitFunError::tool(CoreServiceAgentRuntime::runtime_error_message(error)) })?; - let current_session_id = self.current_workspace_session(context, &workspace); - let result_for_assistant = - self.build_list_result_for_assistant(&workspace, &sessions, current_session_id); + let current_session_id = + self.current_workspace_session(context, &workspace.display_workspace); + let result_for_assistant = self.build_list_result_for_assistant( + &workspace.display_workspace, + &sessions, + current_session_id, + ); Ok(vec![ToolResult::Result { data: json!({ "success": true, "action": "list", - "workspace": workspace.clone(), + "workspace": workspace.display_workspace.clone(), "current_session_id": current_session_id, "count": sessions.len(), "sessions": sessions, diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs index 6f033500b..6296f684b 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs @@ -251,19 +251,21 @@ Examples: let coordinator = get_global_coordinator() .ok_or_else(|| BitFunError::tool("coordinator not initialized".to_string()))?; let workspace = coordinator - .resolve_session_workspace_path(&session_id) + .get_session_manager() + .resolve_session_workspace_binding(&session_id) .await - .map(|path| path.to_string_lossy().to_string()) .ok_or_else(|| { BitFunError::NotFound(format!( "Workspace for session '{}' could not be resolved", session_id )) })?; + let display_workspace = workspace.root_path_string(); + let session_storage_path = workspace.session_storage_path(); let manager = PersistenceManager::new(Arc::new(PathManager::new()?))?; let transcript = manager .export_session_transcript( - std::path::Path::new(&workspace), + &session_storage_path, &session_id, &SessionTranscriptExportOptions { tools: params.tools.unwrap_or(false), @@ -277,7 +279,7 @@ Examples: Ok(vec![ToolResult::Result { data: json!({ "success": true, - "workspace": workspace, + "workspace": display_workspace, "transcript": transcript, }), result_for_assistant: Some(format!( diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/session_message_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/session_message_tool.rs index d28fde7a2..cd20ec679 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/session_message_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/session_message_tool.rs @@ -11,7 +11,7 @@ use crate::util::errors::{BitFunError, BitFunResult}; use async_trait::async_trait; use bitfun_runtime_ports::{ AgentDialogPrependedReminder, AgentDialogTurnRequest, AgentSessionCreateRequest, - AgentSessionListRequest, AgentSessionReplyRoute, AgentSessionWorkspaceRequest, + AgentSessionReplyRoute, AgentSessionWorkspaceBinding, AgentSessionWorkspaceRequest, }; use serde::Deserialize; use serde_json::{json, Value}; @@ -20,6 +20,13 @@ use std::path::Path; /// SessionMessage tool - send a message to another session via the dialog scheduler pub struct SessionMessageTool; +#[derive(Debug, Clone)] +struct SessionMessageWorkspaceTarget { + workspace_path: String, + remote_connection_id: Option, + remote_ssh_host: Option, +} + impl Default for SessionMessageTool { fn default() -> Self { Self::new() @@ -161,6 +168,39 @@ impl SessionMessageTool { Ok(format!("session-{}", creator_session_id)) } + fn workspace_target_from_context( + &self, + workspace_path: String, + context: &ToolUseContext, + ) -> SessionMessageWorkspaceTarget { + let remote_connection_id = context + .workspace + .as_ref() + .and_then(|workspace| workspace.connection_id().map(ToOwned::to_owned)); + let remote_ssh_host = context + .workspace + .as_ref() + .filter(|workspace| workspace.is_remote()) + .map(|workspace| workspace.session_identity.hostname.clone()) + .filter(|value| !value.trim().is_empty()); + SessionMessageWorkspaceTarget { + workspace_path, + remote_connection_id, + remote_ssh_host, + } + } + + fn workspace_target_from_binding( + &self, + binding: AgentSessionWorkspaceBinding, + ) -> SessionMessageWorkspaceTarget { + SessionMessageWorkspaceTarget { + workspace_path: binding.workspace_path, + remote_connection_id: binding.remote_connection_id, + remote_ssh_host: binding.remote_ssh_host, + } + } + fn format_forwarded_message( &self, message: &str, @@ -445,6 +485,16 @@ Allowed agent types when creating a session: .map_err(|e| BitFunError::tool(format!("Invalid input: {}", e)))?; let source_session_id = self.sender_session_id(context)?.to_string(); let source_workspace = self.sender_workspace(context)?; + let source_remote_connection_id = context + .workspace + .as_ref() + .and_then(|workspace| workspace.connection_id().map(ToOwned::to_owned)); + let source_remote_ssh_host = context + .workspace + .as_ref() + .filter(|workspace| workspace.is_remote()) + .map(|workspace| workspace.session_identity.hostname.clone()) + .filter(|value| !value.trim().is_empty()); let coordinator = get_global_coordinator() .ok_or_else(|| BitFunError::tool("coordinator not initialized".to_string()))?; @@ -456,7 +506,7 @@ Allowed agent types when creating a session: ) .map_err(BitFunError::tool)?; - let (target_session_id, target_agent_type, created_session_id, workspace) = + let (target_session_id, target_agent_type, created_session_id, workspace_target) = if let Some(target_session_id) = params.session_id.clone() { if source_session_id == target_session_id { return Err(BitFunError::tool( @@ -464,50 +514,54 @@ Allowed agent types when creating a session: )); } - let workspace = if let Some(workspace) = params.workspace.as_deref() { - self.resolve_workspace(workspace, context)? - } else { - runtime - .resolve_session_workspace_path(AgentSessionWorkspaceRequest { - session_id: target_session_id.clone(), - }) - .await - .map_err(|error| { - BitFunError::tool(CoreServiceAgentRuntime::runtime_error_message(error)) - })? - .ok_or_else(|| { - BitFunError::NotFound(format!( - "Workspace for session '{}' could not be resolved", - target_session_id - )) - })? - }; - let existing_sessions = runtime - .list_sessions(AgentSessionListRequest { - workspace_path: workspace.clone(), + let workspace_target = runtime + .resolve_session_workspace_binding(AgentSessionWorkspaceRequest { + session_id: target_session_id.clone(), }) .await .map_err(|error| { BitFunError::tool(CoreServiceAgentRuntime::runtime_error_message(error)) })?; - let target_session = existing_sessions - .iter() - .find(|session| session.session_id == target_session_id.as_str()) - .ok_or_else(|| { - BitFunError::NotFound(format!( + let workspace_target = workspace_target.ok_or_else(|| { + BitFunError::NotFound(format!( + "Workspace for session '{}' could not be resolved", + target_session_id + )) + })?; + let workspace_target = self.workspace_target_from_binding(workspace_target); + + if let Some(workspace) = params.workspace.as_deref() { + let requested_workspace = self.resolve_workspace(workspace, context)?; + let requested_target = + self.workspace_target_from_context(requested_workspace.clone(), context); + let remote_mismatch = requested_target + .remote_connection_id + .as_deref() + .zip(workspace_target.remote_connection_id.as_deref()) + .map(|(left, right)| left != right) + .unwrap_or(false); + if requested_target.workspace_path != workspace_target.workspace_path + || remote_mismatch + { + return Err(BitFunError::NotFound(format!( "Session '{}' not found in workspace '{}'", - target_session_id, workspace - )) - })?; + target_session_id, requested_target.workspace_path + ))); + } + } - let persisted_agent_type = target_session.agent_type.trim(); - let target_agent_type = if persisted_agent_type.is_empty() { - "agentic".to_string() - } else { - persisted_agent_type.to_string() - }; + let target_agent_type = runtime + .resolve_session_agent_type(&target_session_id) + .await + .map_err(|error| { + BitFunError::tool(CoreServiceAgentRuntime::runtime_error_message(error)) + })? + .filter(|value| !value.trim().is_empty()) + .ok_or_else(|| { + BitFunError::NotFound(format!("Session '{}' not found", target_session_id)) + })?; - (target_session_id, target_agent_type, None, workspace) + (target_session_id, target_agent_type, None, workspace_target) } else { let workspace = self.resolve_workspace( params.workspace.as_deref().ok_or_else(|| { @@ -517,6 +571,7 @@ Allowed agent types when creating a session: })?, context, )?; + let workspace_target = self.workspace_target_from_context(workspace, context); let session_name = params .session_name .clone() @@ -543,7 +598,9 @@ Allowed agent types when creating a session: .create_session(AgentSessionCreateRequest { session_name, agent_type: agent_type.clone(), - workspace_path: Some(workspace.clone()), + workspace_path: Some(workspace_target.workspace_path.clone()), + remote_connection_id: workspace_target.remote_connection_id.clone(), + remote_ssh_host: workspace_target.remote_ssh_host.clone(), metadata, }) .await @@ -555,7 +612,7 @@ Allowed agent types when creating a session: session.session_id.clone(), session.agent_type.clone(), Some(session.session_id), - workspace, + workspace_target, ) }; @@ -569,11 +626,15 @@ Allowed agent types when creating a session: original_message: Some(params.message.clone()), turn_id: None, agent_type: target_agent_type.clone(), - workspace_path: Some(workspace.clone()), + workspace_path: Some(workspace_target.workspace_path.clone()), + remote_connection_id: workspace_target.remote_connection_id.clone(), + remote_ssh_host: workspace_target.remote_ssh_host.clone(), policy: DialogSubmissionPolicy::for_source(DialogTriggerSource::AgentSession), reply_route: Some(AgentSessionReplyRoute { source_session_id, source_workspace_path: source_workspace, + source_remote_connection_id, + source_remote_ssh_host, }), prepended_reminders: prepended_messages, attachments: Vec::new(), @@ -587,7 +648,7 @@ Allowed agent types when creating a session: Ok(vec![ToolResult::Result { data: json!({ "success": true, - "target_workspace": workspace.clone(), + "target_workspace": workspace_target.workspace_path.clone(), "target_session_id": target_session_id.clone(), "target_agent_type": target_agent_type.clone(), "created_session_id": created_session_id.clone(), @@ -595,12 +656,12 @@ Allowed agent types when creating a session: result_for_assistant: Some(if let Some(created_session_id) = created_session_id { format!( "Created session '{}' and accepted the message in workspace '{}' using agent type '{}'.", - created_session_id, workspace, target_agent_type + created_session_id, workspace_target.workspace_path, target_agent_type ) } else { format!( "Message accepted for session '{}' in workspace '{}' using agent type '{}'.", - target_session_id, workspace, target_agent_type + target_session_id, workspace_target.workspace_path, target_agent_type ) }), image_attachments: None, diff --git a/src/crates/assembly/core/src/service/cron/service.rs b/src/crates/assembly/core/src/service/cron/service.rs index 4025c6049..78ffb3ed1 100644 --- a/src/crates/assembly/core/src/service/cron/service.rs +++ b/src/crates/assembly/core/src/service/cron/service.rs @@ -13,6 +13,7 @@ use crate::agentic::coordination::{ DialogTriggerSource, }; use crate::agentic::core::SessionConfig; +use crate::agentic::workspace::WorkspaceBinding; use crate::infrastructure::PathManager; use crate::service_agent_runtime::CoreServiceAgentRuntime; use crate::util::errors::{BitFunError, BitFunResult}; @@ -119,17 +120,21 @@ impl CronService { let jobs = self.jobs.read().await; jobs.values() .filter(|job| { - matches_workspace_filter( + let workspace_matches = matches_workspace_filter( job.workspace(), workspace_path, workspace_id, remote_connection_id, - ) && session_id + ); + let session_matches = session_id .map(|session_id| job.session_id() == Some(session_id)) - .unwrap_or(true) - && target_kind - .map(|target_kind| job.target_kind() == target_kind) - .unwrap_or(true) + .unwrap_or(true); + let target_matches = target_kind + .map(|target_kind| job.target_kind() == target_kind) + .unwrap_or(true); + let included = workspace_matches && session_matches && target_matches; + + included }) .cloned() .collect::>() @@ -140,11 +145,12 @@ impl CronService { } pub async fn create_job(&self, request: CreateCronJobRequest) -> BitFunResult { + let target = self.canonicalize_target(request.target).await?; let _guard = self.mutation_lock.lock().await; let mut jobs = self.jobs.write().await; let current_ms = now_ms(); let schedule = materialize_schedule(request.schedule, current_ms); - let target = materialize_target(request.target); + validate_request_fields(&request.name, &request.payload, &target)?; validate_schedule(&schedule, current_ms)?; @@ -166,6 +172,7 @@ impl CronService { } jobs.insert(job.id.clone(), job.clone()); + self.persist_jobs_locked(&jobs).await?; drop(jobs); self.wakeup.notify_one(); @@ -178,6 +185,11 @@ impl CronService { job_id: &str, request: UpdateCronJobRequest, ) -> BitFunResult { + let canonicalized_target = match request.target { + Some(target) => Some(self.canonicalize_target(target).await?), + None => None, + }; + let _guard = self.mutation_lock.lock().await; let mut jobs = self.jobs.write().await; let current_ms = now_ms(); @@ -191,8 +203,8 @@ impl CronService { if let Some(payload) = request.payload { job.payload = payload; } - if let Some(target) = request.target { - job.target = materialize_target(target); + if let Some(target) = canonicalized_target { + job.target = target; } if let Some(enabled) = request.enabled { job.enabled = enabled; @@ -560,6 +572,8 @@ impl CronService { turn_id: Some(enqueue_input.turn_id.clone()), agent_type: resolved.agent_type, workspace_path: Some(resolved.workspace_path), + remote_connection_id: resolved.remote_connection_id, + remote_ssh_host: resolved.remote_ssh_host, policy: scheduled_job_policy(), reply_route: None, prepended_reminders: enqueue_input.prepended_messages.clone(), @@ -589,6 +603,8 @@ impl CronService { Ok(ResolvedEnqueueSubmission { session_id: session_id.clone(), workspace_path: workspace.workspace_path.clone(), + remote_connection_id: workspace.remote_connection_id.clone(), + remote_ssh_host: workspace.remote_ssh_host.clone(), agent_type, }) } @@ -620,12 +636,47 @@ impl CronService { Ok(ResolvedEnqueueSubmission { session_id: created.session_id, workspace_path: workspace.workspace_path.clone(), + remote_connection_id: workspace.remote_connection_id.clone(), + remote_ssh_host: workspace.remote_ssh_host.clone(), agent_type: created.agent_type, }) } } } + async fn canonicalize_target(&self, target: CronJobTarget) -> BitFunResult { + let mut target = materialize_target(target); + + if let CronJobTarget::Session { + session_id, + workspace, + } = &mut target + { + *workspace = + Self::resolve_session_target_workspace_ref(&self.coordinator, session_id).await?; + } + + Ok(target) + } + + async fn resolve_session_target_workspace_ref( + coordinator: &ConversationCoordinator, + session_id: &str, + ) -> BitFunResult { + let binding = coordinator + .get_session_manager() + .resolve_session_workspace_binding(session_id) + .await + .ok_or_else(|| { + BitFunError::validation(format!( + "Unable to resolve workspace for session '{}'", + session_id + )) + })?; + + Ok(workspace_ref_from_binding(&binding)) + } + async fn persist_jobs_locked(&self, jobs: &HashMap) -> BitFunResult<()> { self.store .save_jobs(jobs.values().cloned().collect::>()) @@ -732,7 +783,7 @@ fn materialize_workspace_ref(workspace: CronWorkspaceRef) -> CronWorkspaceRef { .workspace_id .map(|value| value.trim().to_string()) .filter(|value| !value.is_empty()), - workspace_path: workspace.workspace_path.trim().to_string(), + workspace_path: normalize_workspace_path_for_matching(&workspace.workspace_path), remote_connection_id: workspace .remote_connection_id .map(|value| value.trim().to_string()) @@ -744,6 +795,19 @@ fn materialize_workspace_ref(workspace: CronWorkspaceRef) -> CronWorkspaceRef { } } +fn workspace_ref_from_binding(binding: &WorkspaceBinding) -> CronWorkspaceRef { + CronWorkspaceRef { + workspace_id: binding.workspace_id.clone(), + workspace_path: normalize_workspace_path_for_matching(&binding.root_path_string()), + remote_connection_id: binding.connection_id().map(ToOwned::to_owned), + remote_ssh_host: if binding.is_remote() { + Some(binding.session_identity.hostname.clone()).filter(|value| !value.trim().is_empty()) + } else { + None + }, + } +} + fn materialize_launch_spec(launch: CronLaunchSpec) -> CronLaunchSpec { CronLaunchSpec { agent_type: normalize_agent_type(&launch.agent_type), @@ -800,8 +864,10 @@ fn matches_workspace_filter( workspace_id: Option<&str>, remote_connection_id: Option<&str>, ) -> bool { + let normalized_job_workspace_path = + normalize_workspace_path_for_matching(&workspace.workspace_path); let workspace_path_matches = workspace_path - .map(|value| workspace.workspace_path == value) + .map(|value| normalized_job_workspace_path == normalize_workspace_path_for_matching(value)) .unwrap_or(true); let workspace_id_matches = workspace_id .map(|value| { @@ -815,6 +881,50 @@ fn matches_workspace_filter( workspace_path_matches && workspace_id_matches && remote_connection_matches } +fn normalize_workspace_path_for_matching(path: &str) -> String { + let mut normalized = path.trim().replace('\\', "/"); + + if normalized.starts_with("file://") { + normalized = normalized.trim_start_matches("file://").to_string(); + } + + if normalized.len() >= 4 + && normalized.starts_with('/') + && normalized.as_bytes()[2] == b':' + && normalized.as_bytes()[1].is_ascii_alphabetic() + { + normalized = normalized.trim_start_matches('/').to_string(); + } + + while normalized.contains("//") { + normalized = normalized.replace("//", "/"); + } + + if normalized.len() >= 2 + && normalized.as_bytes()[1] == b':' + && normalized.as_bytes()[0].is_ascii_alphabetic() + { + normalized = format!( + "{}{}", + normalized[..1].to_ascii_uppercase(), + &normalized[1..] + ); + } + + if normalized != "/" && !is_windows_drive_root(&normalized) { + normalized = normalized.trim_end_matches('/').to_string(); + } + + normalized +} + +fn is_windows_drive_root(path: &str) -> bool { + path.len() == 3 + && path.as_bytes()[1] == b':' + && path.as_bytes()[2] == b'/' + && path.as_bytes()[0].is_ascii_alphabetic() +} + fn next_wakeup_for_job(job: &CronJob) -> Option { job.state.next_wakeup_at_ms() } @@ -865,6 +975,8 @@ struct EnqueueInput { struct ResolvedEnqueueSubmission { session_id: String, workspace_path: String, + remote_connection_id: Option, + remote_ssh_host: Option, agent_type: String, } @@ -879,3 +991,57 @@ fn submit_target_session_id(enqueue_input: &EnqueueInput) -> &str { fn cron_enqueue_error_is_missing_session(error: &str) -> bool { error.contains("Session metadata not found") } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn materialize_workspace_ref_normalizes_windows_style_paths() { + let workspace = materialize_workspace_ref(CronWorkspaceRef { + workspace_id: None, + workspace_path: r"c:\Users\wsp\.bitfun\personal_assistant\workspace\".to_string(), + remote_connection_id: None, + remote_ssh_host: None, + }); + + assert_eq!( + workspace.workspace_path, + "C:/Users/wsp/.bitfun/personal_assistant/workspace" + ); + } + + #[test] + fn matches_workspace_filter_tolerates_separator_differences() { + let workspace = CronWorkspaceRef { + workspace_id: Some("local_workspace".to_string()), + workspace_path: r"C:\Users\wsp\.bitfun\personal_assistant\workspace".to_string(), + remote_connection_id: None, + remote_ssh_host: None, + }; + + assert!(matches_workspace_filter( + &workspace, + Some("C:/Users/wsp/.bitfun/personal_assistant/workspace"), + Some("local_workspace"), + None, + )); + } + + #[test] + fn matches_workspace_filter_normalizes_remote_like_paths() { + let workspace = CronWorkspaceRef { + workspace_id: None, + workspace_path: "/home/wsp/projects/test/".to_string(), + remote_connection_id: Some("ssh-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), + }; + + assert!(matches_workspace_filter( + &workspace, + Some(r"\home\wsp\projects\test"), + None, + Some("ssh-1"), + )); + } +} diff --git a/src/crates/assembly/core/src/service_agent_runtime.rs b/src/crates/assembly/core/src/service_agent_runtime.rs index 00ed15960..b776b21f0 100644 --- a/src/crates/assembly/core/src/service_agent_runtime.rs +++ b/src/crates/assembly/core/src/service_agent_runtime.rs @@ -1026,6 +1026,8 @@ impl RemoteDialogRuntimeHost for CoreRemoteDialogRuntimeHost<'_> { turn_id: Some(submission.turn_id), agent_type: submission.resolved_agent_type, workspace_path: submission.binding_workspace, + remote_connection_id: None, + remote_ssh_host: None, policy, reply_route: None, prepended_reminders: Vec::new(), diff --git a/src/crates/contracts/runtime-ports/src/lib.rs b/src/crates/contracts/runtime-ports/src/lib.rs index a21ccb87c..56c2638f6 100644 --- a/src/crates/contracts/runtime-ports/src/lib.rs +++ b/src/crates/contracts/runtime-ports/src/lib.rs @@ -577,6 +577,10 @@ pub struct AgentSessionCreateRequest { pub agent_type: String, #[serde(skip_serializing_if = "Option::is_none")] pub workspace_path: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_connection_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_ssh_host: Option, #[serde(default, skip_serializing_if = "serde_json::Map::is_empty")] pub metadata: serde_json::Map, } @@ -594,6 +598,10 @@ pub struct AgentSessionCreateResult { #[serde(rename_all = "camelCase")] pub struct AgentSessionListRequest { pub workspace_path: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_connection_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_ssh_host: Option, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -611,6 +619,10 @@ pub struct AgentSessionSummary { pub struct AgentSessionDeleteRequest { pub workspace_path: String, pub session_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_connection_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_ssh_host: Option, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -619,6 +631,16 @@ pub struct AgentSessionWorkspaceRequest { pub session_id: String, } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentSessionWorkspaceBinding { + pub workspace_path: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_connection_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_ssh_host: Option, +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct AgentSubmissionRequest { @@ -646,6 +668,10 @@ pub struct AgentDialogTurnRequest { pub agent_type: String, #[serde(skip_serializing_if = "Option::is_none")] pub workspace_path: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_connection_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_ssh_host: Option, pub policy: DialogSubmissionPolicy, #[serde(default, skip_serializing_if = "Option::is_none")] pub reply_route: Option, @@ -841,6 +867,10 @@ pub const fn should_skip_agent_session_reply( pub struct AgentSessionReplyRoute { pub source_session_id: String, pub source_workspace_path: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub source_remote_connection_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub source_remote_ssh_host: Option, } /// Outcome for steering a message into an already-running dialog turn. @@ -1199,6 +1229,11 @@ pub trait AgentSessionManagementPort: Send + Sync { &self, request: AgentSessionWorkspaceRequest, ) -> PortResult>; + + async fn resolve_session_workspace_binding( + &self, + request: AgentSessionWorkspaceRequest, + ) -> PortResult>; } #[async_trait::async_trait] @@ -1619,10 +1654,14 @@ mod tests { let route = AgentSessionReplyRoute { source_session_id: "requester_session".to_string(), source_workspace_path: "/workspace/requester".to_string(), + source_remote_connection_id: Some("conn-1".to_string()), + source_remote_ssh_host: Some("host-1".to_string()), }; assert_eq!(route.source_session_id, "requester_session"); assert_eq!(route.source_workspace_path, "/workspace/requester"); + assert_eq!(route.source_remote_connection_id.as_deref(), Some("conn-1")); + assert_eq!(route.source_remote_ssh_host.as_deref(), Some("host-1")); } #[test] @@ -1902,6 +1941,8 @@ mod tests { turn_id: Some("turn_1".to_string()), agent_type: "agentic".to_string(), workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), policy: DialogSubmissionPolicy::new( AgentSubmissionSource::RemoteRelay, DialogQueuePriority::High, @@ -1910,6 +1951,8 @@ mod tests { reply_route: Some(AgentSessionReplyRoute { source_session_id: "source_session".to_string(), source_workspace_path: "/workspace/source".to_string(), + source_remote_connection_id: Some("conn-1".to_string()), + source_remote_ssh_host: Some("host-1".to_string()), }), prepended_reminders: vec![AgentDialogPrependedReminder { kind: "session_message_request".to_string(), @@ -1931,10 +1974,14 @@ mod tests { assert_eq!(json["turnId"], "turn_1"); assert_eq!(json["agentType"], "agentic"); assert_eq!(json["workspacePath"], "/workspace/project"); + assert_eq!(json["remoteConnectionId"], "conn-1"); + assert_eq!(json["remoteSshHost"], "host-1"); assert_eq!(json["policy"]["triggerSource"], "remote_relay"); assert_eq!(json["policy"]["queuePriority"], "high"); assert_eq!(json["policy"]["skipToolConfirmation"], true); assert_eq!(json["replyRoute"]["sourceSessionId"], "source_session"); + assert_eq!(json["replyRoute"]["sourceRemoteConnectionId"], "conn-1"); + assert_eq!(json["replyRoute"]["sourceRemoteSshHost"], "host-1"); assert_eq!( json["prependedReminders"][0]["kind"], "session_message_request" @@ -2023,6 +2070,8 @@ mod tests { fn agent_session_management_contracts_serialize_stable_shape() { let list_request = AgentSessionListRequest { workspace_path: "/workspace/project".to_string(), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), }; let summary = AgentSessionSummary { session_id: "session_1".to_string(), @@ -2034,23 +2083,39 @@ mod tests { let delete_request = AgentSessionDeleteRequest { workspace_path: "/workspace/project".to_string(), session_id: "session_1".to_string(), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), }; let workspace_request = AgentSessionWorkspaceRequest { session_id: "session_1".to_string(), }; + let workspace_binding = AgentSessionWorkspaceBinding { + workspace_path: "/workspace/project".to_string(), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), + }; let list_json = serde_json::to_value(list_request).expect("serialize list request"); let summary_json = serde_json::to_value(summary).expect("serialize summary"); let delete_json = serde_json::to_value(delete_request).expect("serialize delete request"); let workspace_json = serde_json::to_value(workspace_request).expect("serialize workspace request"); + let binding_json = + serde_json::to_value(workspace_binding).expect("serialize workspace binding"); assert_eq!(list_json["workspacePath"], "/workspace/project"); + assert_eq!(list_json["remoteConnectionId"], "conn-1"); + assert_eq!(list_json["remoteSshHost"], "host-1"); assert_eq!(summary_json["sessionId"], "session_1"); assert_eq!(summary_json["createdAtMs"], 1000); assert_eq!(summary_json["lastActiveAtMs"], 2000); assert_eq!(delete_json["sessionId"], "session_1"); + assert_eq!(delete_json["remoteConnectionId"], "conn-1"); + assert_eq!(delete_json["remoteSshHost"], "host-1"); assert_eq!(workspace_json["sessionId"], "session_1"); + assert_eq!(binding_json["workspacePath"], "/workspace/project"); + assert_eq!(binding_json["remoteConnectionId"], "conn-1"); + assert_eq!(binding_json["remoteSshHost"], "host-1"); } #[test] diff --git a/src/crates/execution/agent-runtime/src/runtime.rs b/src/crates/execution/agent-runtime/src/runtime.rs index 62f7a8b3f..3a12f4b91 100644 --- a/src/crates/execution/agent-runtime/src/runtime.rs +++ b/src/crates/execution/agent-runtime/src/runtime.rs @@ -13,10 +13,11 @@ use bitfun_runtime_ports::{ AgentBackgroundResultRequest, AgentDialogTurnPort, AgentDialogTurnRequest, AgentInputAttachment, AgentLifecycleDeliveryPort, AgentSessionCreateRequest, AgentSessionCreateResult, AgentSessionDeleteRequest, AgentSessionListRequest, - AgentSessionManagementPort, AgentSessionSummary, AgentSessionWorkspaceRequest, - AgentSubmissionPort, AgentSubmissionRequest, AgentSubmissionResult, AgentSubmissionSource, - AgentThreadGoalDeliveryRequest, AgentTurnCancellationPort, AgentTurnCancellationRequest, - AgentTurnCancellationResult, DialogSubmitOutcome, PortError, RuntimeEventEnvelope, + AgentSessionManagementPort, AgentSessionSummary, AgentSessionWorkspaceBinding, + AgentSessionWorkspaceRequest, AgentSubmissionPort, AgentSubmissionRequest, + AgentSubmissionResult, AgentSubmissionSource, AgentThreadGoalDeliveryRequest, + AgentTurnCancellationPort, AgentTurnCancellationRequest, AgentTurnCancellationResult, + DialogSubmitOutcome, PortError, RuntimeEventEnvelope, }; use bitfun_runtime_services::RuntimeServices; @@ -458,6 +459,20 @@ impl AgentRuntime { .map_err(RuntimeError::from) } + pub async fn resolve_session_workspace_binding( + &self, + request: AgentSessionWorkspaceRequest, + ) -> Result, RuntimeError> { + let session_management = self + .session_management + .as_ref() + .ok_or(RuntimeError::MissingSessionManagementPort)?; + session_management + .resolve_session_workspace_binding(request) + .await + .map_err(RuntimeError::from) + } + pub async fn submit_turn( &self, request: AgentSubmissionRequest, @@ -569,6 +584,8 @@ impl AgentRuntime { session_name, agent_type, workspace_path, + remote_connection_id: None, + remote_ssh_host: None, metadata, }) .await?; @@ -621,6 +638,7 @@ mod tests { listed_sessions: Mutex>, deleted_sessions: Mutex>, workspace_requests: Mutex>, + workspace_binding_requests: Mutex>, resolved_agent_type: Option, } @@ -652,6 +670,21 @@ mod tests { self.workspace_requests.lock().unwrap().push(request); Ok(Some("/workspace/project".to_string())) } + + async fn resolve_session_workspace_binding( + &self, + request: AgentSessionWorkspaceRequest, + ) -> PortResult> { + self.workspace_binding_requests + .lock() + .unwrap() + .push(request); + Ok(Some(AgentSessionWorkspaceBinding { + workspace_path: "/workspace/project".to_string(), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), + })) + } } #[async_trait::async_trait] @@ -889,6 +922,8 @@ mod tests { let err = runtime .list_sessions(AgentSessionListRequest { workspace_path: "/workspace/project".to_string(), + remote_connection_id: None, + remote_ssh_host: None, }) .await .unwrap_err(); @@ -908,6 +943,8 @@ mod tests { let sessions = runtime .list_sessions(AgentSessionListRequest { workspace_path: "/workspace/project".to_string(), + remote_connection_id: None, + remote_ssh_host: None, }) .await .expect("list sessions"); @@ -915,6 +952,8 @@ mod tests { .delete_session(AgentSessionDeleteRequest { workspace_path: "/workspace/project".to_string(), session_id: "session_1".to_string(), + remote_connection_id: None, + remote_ssh_host: None, }) .await .expect("delete session"); @@ -924,12 +963,25 @@ mod tests { }) .await .expect("resolve workspace"); + let workspace_binding = runtime + .resolve_session_workspace_binding(AgentSessionWorkspaceRequest { + session_id: "session_1".to_string(), + }) + .await + .expect("resolve workspace binding") + .expect("workspace binding"); assert_eq!(sessions[0].session_id, "session_1"); assert_eq!(workspace_path.as_deref(), Some("/workspace/project")); + assert_eq!(workspace_binding.workspace_path, "/workspace/project"); + assert_eq!( + workspace_binding.remote_connection_id.as_deref(), + Some("conn-1") + ); assert_eq!(ports.listed_sessions.lock().unwrap().len(), 1); assert_eq!(ports.deleted_sessions.lock().unwrap().len(), 1); assert_eq!(ports.workspace_requests.lock().unwrap().len(), 1); + assert_eq!(ports.workspace_binding_requests.lock().unwrap().len(), 1); } #[tokio::test] @@ -948,6 +1000,8 @@ mod tests { turn_id: Some("turn_1".to_string()), agent_type: "agentic".to_string(), workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: None, + remote_ssh_host: None, policy: DialogSubmissionPolicy::new( AgentSubmissionSource::RemoteRelay, DialogQueuePriority::Normal, @@ -1001,6 +1055,8 @@ mod tests { turn_id: Some("turn_1".to_string()), agent_type: "agentic".to_string(), workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: None, + remote_ssh_host: None, policy: DialogSubmissionPolicy::new( AgentSubmissionSource::RemoteRelay, DialogQueuePriority::High, diff --git a/src/crates/execution/agent-runtime/src/scheduler.rs b/src/crates/execution/agent-runtime/src/scheduler.rs index aa6f11905..4196c9f41 100644 --- a/src/crates/execution/agent-runtime/src/scheduler.rs +++ b/src/crates/execution/agent-runtime/src/scheduler.rs @@ -292,6 +292,8 @@ impl DialogTurnQueue { pub struct AgentSessionReplyPlan { pub target_session_id: String, pub target_workspace_path: String, + pub target_remote_connection_id: Option, + pub target_remote_ssh_host: Option, pub user_input: String, pub reminder_text: String, } @@ -828,6 +830,8 @@ pub fn resolve_agent_session_reply_action( AgentSessionReplyAction::Forward(AgentSessionReplyPlan { target_session_id: reply_route.source_session_id.clone(), target_workspace_path: reply_route.source_workspace_path.clone(), + target_remote_connection_id: reply_route.source_remote_connection_id.clone(), + target_remote_ssh_host: reply_route.source_remote_ssh_host.clone(), user_input: outcome.reply_text(), reminder_text: format!( "This message is an automated reply to a previous SessionMessage call, not a human user message.\n\ diff --git a/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs b/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs index 132363417..9afc921d0 100644 --- a/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs +++ b/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs @@ -300,6 +300,8 @@ fn active_dialog_turn_owns_agent_session_reply_suppression_facts() { let route = AgentSessionReplyRoute { source_session_id: "source-session".to_string(), source_workspace_path: "workspace".to_string(), + source_remote_connection_id: Some("conn-1".to_string()), + source_remote_ssh_host: Some("host-1".to_string()), }; let turn = ActiveDialogTurn::new( "turn-1".to_string(), @@ -400,6 +402,8 @@ fn agent_session_reply_action_forwards_completed_outcome_with_legacy_reminder_te }; assert_eq!(plan.target_session_id, "source-session"); assert_eq!(plan.target_workspace_path, "workspace"); + assert_eq!(plan.target_remote_connection_id.as_deref(), Some("conn-1")); + assert_eq!(plan.target_remote_ssh_host.as_deref(), Some("host-1")); assert_eq!(plan.user_input, "done"); assert_eq!( plan.reminder_text, @@ -609,6 +613,8 @@ fn agent_session_turn(source_session_id: &str) -> ActiveDialogTurn { Some(AgentSessionReplyRoute { source_session_id: source_session_id.to_string(), source_workspace_path: "workspace".to_string(), + source_remote_connection_id: Some("conn-1".to_string()), + source_remote_ssh_host: Some("host-1".to_string()), }), ) } diff --git a/src/crates/interfaces/acp/src/runtime/prompt.rs b/src/crates/interfaces/acp/src/runtime/prompt.rs index 5700a150c..4c8428418 100644 --- a/src/crates/interfaces/acp/src/runtime/prompt.rs +++ b/src/crates/interfaces/acp/src/runtime/prompt.rs @@ -51,6 +51,8 @@ impl BitfunAcpRuntime { None, acp_session.mode_id.clone(), Some(acp_session.cwd.clone()), + None, + None, DialogSubmissionPolicy::for_source(DialogTriggerSource::Cli), Some(acp_user_message_metadata()), ) @@ -67,6 +69,8 @@ impl BitfunAcpRuntime { None, acp_session.mode_id.clone(), Some(acp_session.cwd.clone()), + None, + None, DialogSubmissionPolicy::for_source(DialogTriggerSource::Cli), Some(acp_user_message_metadata()), ) diff --git a/src/crates/services/services-integrations/src/remote_connect.rs b/src/crates/services/services-integrations/src/remote_connect.rs index 5942e14fb..a5cad2e7f 100644 --- a/src/crates/services/services-integrations/src/remote_connect.rs +++ b/src/crates/services/services-integrations/src/remote_connect.rs @@ -79,6 +79,8 @@ pub fn build_remote_session_create_request( session_name: session_name.into(), agent_type: agent_type.into(), workspace_path: workspace_path.map(Into::into), + remote_connection_id: None, + remote_ssh_host: None, metadata, } } From a18e19fc9a999c732ede44eb6a40665a39ef944f Mon Sep 17 00:00:00 2001 From: wsp1911 Date: Fri, 26 Jun 2026 15:07:44 +0800 Subject: [PATCH 2/4] fix: restore remote sessions from resolved storage dirs - Fix remote dialog, cancel, polling, and model lookup paths to restore evicted sessions from the bound session storage directory instead of the logical remote workspace root. - Replace the ambiguous session workspace path resolver with workspace bindings that expose separate logical workspace paths and final on-disk session storage dirs. - Carry remote workspace identity through remote dialog submissions so queued or restored turns keep the correct SSH connection and host binding. - Route empty-context history reloads through the same resolved session storage path used by initial session restore. - Make session storage resolution consistent for local and remote workspaces: local roots resolve to managed project sessions dirs, remote roots resolve to SSH mirror sessions dirs, and unresolved remote sessions stay under the dedicated unresolved tree. - Teach persistence and insight/session listing paths to accept already-resolved sessions dirs without re-slugging them, while rejecting remote runtime roots that are not actual sessions dirs. - Add focused coverage for local, remote, unresolved remote, already-resolved sessions-dir, and remote dialog workspace identity behavior. --- src/apps/cli/src/agent/core_adapter.rs | 4 + src/apps/desktop/src/api/agentic_api.rs | 17 +- .../src/agentic/coordination/coordinator.rs | 299 ++++++++++++------ .../src/agentic/coordination/scheduler.rs | 45 ++- .../core/src/agentic/deep_review/report.rs | 6 +- .../src/agentic/insights/session_paths.rs | 44 +-- .../core/src/agentic/persistence/manager.rs | 105 ++++-- .../src/agentic/session/session_manager.rs | 277 ++++++++++------ .../src/agentic/session/session_store_port.rs | 78 ++++- .../tools/implementations/bash_tool.rs | 20 ++ .../tools/implementations/code_review_tool.rs | 4 +- .../implementations/session_history_tool.rs | 4 +- .../implementations/session_message_tool.rs | 59 +++- .../tools/implementations/task_tool.rs | 4 +- .../src/agentic/tools/tool_context_runtime.rs | 10 + .../src/agentic/tools/tool_result_storage.rs | 8 +- .../assembly/core/src/agentic/workspace.rs | 44 ++- .../infrastructure/app_paths/path_manager.rs | 8 +- .../assembly/core/src/miniapp/builtin/mod.rs | 28 +- .../assembly/core/src/miniapp/manager.rs | 28 +- .../src/product_runtime/runtime_services.rs | 2 +- .../remote_connect/bot/command_router.rs | 1 + .../service/remote_connect/remote_server.rs | 12 +- .../src/service/remote_ssh/workspace_state.rs | 114 +++++-- .../core/src/service/snapshot/manager.rs | 9 + .../core/src/service/workspace_runtime/mod.rs | 2 + .../src/service/workspace_runtime/service.rs | 41 +++ .../core/src/service_agent_runtime.rs | 182 ++++++++--- src/crates/contracts/runtime-ports/src/lib.rs | 56 +++- .../execution/agent-runtime/src/runtime.rs | 49 ++- .../execution/agent-runtime/src/scheduler.rs | 22 ++ .../tests/scheduler_contracts.rs | 10 + .../runtime-services/src/test_support.rs | 2 + .../src/remote_connect.rs | 234 +++++++++++--- .../tests/remote_connect_contracts.rs | 103 +++++- .../startupPerformanceContract.test.ts | 2 +- 36 files changed, 1446 insertions(+), 487 deletions(-) diff --git a/src/apps/cli/src/agent/core_adapter.rs b/src/apps/cli/src/agent/core_adapter.rs index d5f562654..ab102df1e 100644 --- a/src/apps/cli/src/agent/core_adapter.rs +++ b/src/apps/cli/src/agent/core_adapter.rs @@ -230,6 +230,8 @@ impl Agent for CoreAgentAdapter { Some(turn_id.clone()), agent_type.to_string(), Some(self.workspace_path_string()), + None, + None, DialogSubmissionPolicy::for_source(DialogTriggerSource::Cli), None, ) @@ -252,6 +254,8 @@ impl Agent for CoreAgentAdapter { Some(turn_id.clone()), agent_type.to_string(), Some(self.workspace_path_string()), + None, + None, DialogSubmissionPolicy::for_source(DialogTriggerSource::Cli), None, ) diff --git a/src/apps/desktop/src/api/agentic_api.rs b/src/apps/desktop/src/api/agentic_api.rs index e8d2b91e6..20d73178c 100644 --- a/src/apps/desktop/src/api/agentic_api.rs +++ b/src/apps/desktop/src/api/agentic_api.rs @@ -1014,7 +1014,7 @@ async fn ensure_session_for_thread_goal( .ok_or_else(|| format!("Session workspace_path is missing: {session_id}")) } -async fn resolve_session_workspace_path_for_thread_goal_read( +async fn resolve_thread_goal_storage_path( coordinator: &Arc, app_state: &AppState, session_id: &str, @@ -1022,6 +1022,15 @@ async fn resolve_session_workspace_path_for_thread_goal_read( remote_connection_id: Option<&str>, remote_ssh_host: Option<&str>, ) -> Result { + if let Some(storage_path) = coordinator + .get_session_manager() + .resolve_session_workspace_binding(session_id) + .await + .map(|binding| binding.session_storage_dir()) + { + return Ok(storage_path); + } + if let Some(workspace_path) = coordinator .get_session_manager() .get_session(session_id) @@ -1057,7 +1066,7 @@ pub async fn get_session_thread_goal( if session_id.is_empty() { return Err("session_id is required".to_string()); } - let workspace_path = resolve_session_workspace_path_for_thread_goal_read( + let storage_path = resolve_thread_goal_storage_path( coordinator.inner(), app_state.inner(), session_id, @@ -1067,7 +1076,7 @@ pub async fn get_session_thread_goal( ) .await?; let goal = coordinator - .get_thread_goal(session_id, workspace_path.as_path()) + .get_thread_goal(session_id, storage_path.as_path()) .await .map_err(|error| error.to_string())?; Ok(GetSessionThreadGoalResponse { goal }) @@ -1227,6 +1236,8 @@ pub async fn run_init_agents_md( .submit_init_agents_md( session_id.to_string(), workspace_path, + request.remote_connection_id.clone(), + request.remote_ssh_host.clone(), DialogSubmissionPolicy::for_source(DialogTriggerSource::DesktopUi), ) .await diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index 5ca937a9e..ce6d7a889 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -816,6 +816,32 @@ impl ConversationCoordinator { } } + async fn restore_path_for_existing_session(&self, session_id: &str) -> BitFunResult { + if let Some(binding) = self + .session_manager + .resolve_session_workspace_binding(session_id) + .await + { + return Ok(binding.session_storage_dir()); + } + + let session = self + .session_manager + .get_session(session_id) + .ok_or_else(|| BitFunError::NotFound(format!("Session not found: {}", session_id)))?; + session + .config + .workspace_path + .as_deref() + .map(PathBuf::from) + .ok_or_else(|| { + BitFunError::Validation(format!( + "workspace_path is required when restoring session: {}", + session_id + )) + }) + } + async fn is_chinese_locale() -> bool { use crate::service::config::get_global_config_service; use crate::service::config::types::AppConfig; @@ -1677,25 +1703,33 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .await?; if context_messages.is_empty() && !session.dialog_turn_ids.is_empty() { - if let Some(workspace_path) = session.config.workspace_path.as_deref() { - match self - .session_manager - .restore_session(Path::new(workspace_path), session_id) - .await - { - Ok(_) => { - context_messages = self - .session_manager - .get_context_messages(session_id) - .await?; - } - Err(e) => { - debug!( - "Failed to restore parent session context for fork capture: session_id={}, error={}", - session_id, e - ); + match self.restore_path_for_existing_session(session_id).await { + Ok(restore_path) => { + match self + .session_manager + .restore_session(&restore_path, session_id) + .await + { + Ok(_) => { + context_messages = self + .session_manager + .get_context_messages(session_id) + .await?; + } + Err(e) => { + debug!( + "Failed to restore parent session context for fork capture: session_id={}, error={}", + session_id, e + ); + } } } + Err(e) => { + debug!( + "Failed to resolve parent session restore path for fork capture: session_id={}, error={}", + session_id, e + ); + } } } @@ -2100,13 +2134,41 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet }) } + async fn require_main_session_storage_path(&self, session_id: &str) -> BitFunResult { + self.require_main_session_workspace(session_id)?; + self.session_manager + .resolve_session_workspace_binding(session_id) + .await + .map(|binding| binding.session_storage_dir()) + .ok_or_else(|| { + BitFunError::Validation(format!( + "Session storage path is unavailable: {session_id}" + )) + }) + } + + async fn resolve_thread_goal_storage_path( + &self, + session_id: &str, + workspace_path: &Path, + ) -> BitFunResult { + if self.session_manager.get_session(session_id).is_some() { + self.require_main_session_storage_path(session_id).await + } else { + Ok(workspace_path.to_path_buf()) + } + } + pub async fn get_thread_goal( &self, session_id: &str, workspace_path: &Path, ) -> BitFunResult> { + let storage_path = self + .resolve_thread_goal_storage_path(session_id, workspace_path) + .await?; self.thread_goal_store() - .get_thread_goal(session_id, workspace_path) + .get_thread_goal(session_id, storage_path.as_path()) .await } @@ -2115,9 +2177,12 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet session_id: &str, workspace_path: &Path, ) -> BitFunResult<()> { + let storage_path = self + .resolve_thread_goal_storage_path(session_id, workspace_path) + .await?; self.thread_goal_runtime.clear_active_goal(None); self.thread_goal_store() - .clear_thread_goal(session_id, workspace_path) + .clear_thread_goal(session_id, storage_path.as_path()) .await?; self.emit_thread_goal_updated(session_id, None).await; Ok(()) @@ -2126,14 +2191,14 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet pub async fn create_thread_goal( &self, session_id: &str, - workspace_path: &Path, + _workspace_path: &Path, objective: String, token_budget: Option, ) -> BitFunResult { - self.require_main_session_workspace(session_id)?; + let storage_path = self.require_main_session_storage_path(session_id).await?; let goal = self .thread_goal_store() - .create_thread_goal(session_id, workspace_path, objective, token_budget) + .create_thread_goal(session_id, storage_path.as_path(), objective, token_budget) .await?; self.thread_goal_runtime.mark_turn_started("", Some(&goal)); self.emit_thread_goal_updated(session_id, Some(goal.clone())) @@ -2144,12 +2209,13 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet pub async fn update_thread_goal_objective( &self, session_id: &str, - workspace_path: &Path, + _workspace_path: &Path, objective: String, ) -> BitFunResult { - self.require_main_session_workspace(session_id)?; + let storage_path = self.require_main_session_storage_path(session_id).await?; let existing = self - .get_thread_goal(session_id, workspace_path) + .thread_goal_store() + .get_thread_goal(session_id, storage_path.as_path()) .await? .ok_or_else(|| { BitFunError::NotFound(format!( @@ -2166,7 +2232,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .thread_goal_store() .set_thread_goal( session_id, - workspace_path, + storage_path.as_path(), Some(objective), status, None, @@ -2190,12 +2256,15 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet pub async fn set_thread_goal_objective( &self, session_id: &str, - workspace_path: &Path, + _workspace_path: &Path, objective: String, replace_existing: bool, ) -> BitFunResult { - self.require_main_session_workspace(session_id)?; - let previous = self.get_thread_goal(session_id, workspace_path).await?; + let storage_path = self.require_main_session_storage_path(session_id).await?; + let previous = self + .thread_goal_store() + .get_thread_goal(session_id, storage_path.as_path()) + .await?; let status = if previous.is_some() && !replace_existing { None } else { @@ -2205,7 +2274,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .thread_goal_store() .set_thread_goal( session_id, - workspace_path, + storage_path.as_path(), Some(objective), status, None, @@ -2248,6 +2317,16 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .require_main_session_workspace(session_id) .ok() .map(|path| path.to_string_lossy().to_string()); + let (remote_connection_id, remote_ssh_host) = self + .session_manager + .get_session(session_id) + .map(|session| { + ( + session.config.remote_connection_id.clone(), + session.config.remote_ssh_host.clone(), + ) + }) + .unwrap_or((None, None)); let runtime = match CoreServiceAgentRuntime::global_agent_runtime_with_lifecycle_delivery() { Ok(runtime) => runtime, @@ -2264,6 +2343,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet session_id: session_id.to_string(), agent_type, workspace_path, + remote_connection_id, + remote_ssh_host, kind: AgentThreadGoalDeliveryKind::ObjectiveUpdated, goal: goal.clone(), }) @@ -2285,12 +2366,13 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet if !is_usage_limit_error(error) { return; } - let workspace_path = match self.require_main_session_workspace(session_id) { + let storage_path = match self.require_main_session_storage_path(session_id).await { Ok(path) => path, Err(_) => return, }; let Ok(Some(goal)) = self - .get_thread_goal(session_id, workspace_path.as_path()) + .thread_goal_store() + .get_thread_goal(session_id, storage_path.as_path()) .await else { return; @@ -2301,7 +2383,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet if let Err(error) = self .set_thread_goal_status( session_id, - workspace_path.as_path(), + storage_path.as_path(), ThreadGoalStatus::UsageLimited, ) .await @@ -2316,18 +2398,28 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet pub async fn set_thread_goal_status( &self, session_id: &str, - workspace_path: &Path, + _workspace_path: &Path, status: ThreadGoalStatus, ) -> BitFunResult { - self.require_main_session_workspace(session_id)?; - let previous = self.get_thread_goal(session_id, workspace_path).await?; + let storage_path = self.require_main_session_storage_path(session_id).await?; + let previous = self + .thread_goal_store() + .get_thread_goal(session_id, storage_path.as_path()) + .await?; let resuming = status == ThreadGoalStatus::Active && previous .as_ref() .is_some_and(|goal| thread_goal_status_is_resumable(goal.status)); let result = self .thread_goal_store() - .set_thread_goal(session_id, workspace_path, None, Some(status), None, false) + .set_thread_goal( + session_id, + storage_path.as_path(), + None, + Some(status), + None, + false, + ) .await?; if !result.goal.is_active() { self.thread_goal_runtime.clear_active_goal(None); @@ -2346,7 +2438,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet /// Pause an active thread goal after the user manually stops a turn so the UI can offer resume. pub async fn pause_thread_goal_after_user_cancel(&self, session_id: &str) { - let workspace_path = match self.require_main_session_workspace(session_id) { + let storage_path = match self.require_main_session_storage_path(session_id).await { Ok(path) => path, Err(error) => { debug!( @@ -2357,7 +2449,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet } }; let Ok(Some(goal)) = self - .get_thread_goal(session_id, workspace_path.as_path()) + .thread_goal_store() + .get_thread_goal(session_id, storage_path.as_path()) .await else { return; @@ -2366,11 +2459,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet return; } if let Err(error) = self - .set_thread_goal_status( - session_id, - workspace_path.as_path(), - ThreadGoalStatus::Paused, - ) + .set_thread_goal_status(session_id, storage_path.as_path(), ThreadGoalStatus::Paused) .await { warn!( @@ -2404,6 +2493,16 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .require_main_session_workspace(session_id) .ok() .map(|path| path.to_string_lossy().to_string()); + let (remote_connection_id, remote_ssh_host) = self + .session_manager + .get_session(session_id) + .map(|session| { + ( + session.config.remote_connection_id.clone(), + session.config.remote_ssh_host.clone(), + ) + }) + .unwrap_or((None, None)); let session_id = session_id.to_string(); let goal = goal.clone(); tokio::spawn(async move { @@ -2423,6 +2522,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet session_id: session_id.clone(), agent_type, workspace_path, + remote_connection_id, + remote_ssh_host, kind: AgentThreadGoalDeliveryKind::Resumed, goal, }) @@ -2461,9 +2562,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet } async fn load_active_thread_goal(&self, session_id: &str) -> BitFunResult> { - let workspace_path = self.require_main_session_workspace(session_id)?; + let storage_path = self.require_main_session_storage_path(session_id).await?; Ok(self - .get_thread_goal(session_id, workspace_path.as_path()) + .thread_goal_store() + .get_thread_goal(session_id, storage_path.as_path()) .await? .filter(ThreadGoal::is_active)) } @@ -2479,15 +2581,16 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet "Goal objective is required. Use /goal .".to_string(), ) })?; - let workspace_path = self.require_main_session_workspace(&session_id)?; + let storage_path = self.require_main_session_storage_path(&session_id).await?; let existing = self - .get_thread_goal(&session_id, workspace_path.as_path()) + .thread_goal_store() + .get_thread_goal(&session_id, storage_path.as_path()) .await?; let replace_existing = existing.is_some(); let goal = self .set_thread_goal_objective( &session_id, - workspace_path.as_path(), + storage_path.as_path(), objective, replace_existing, ) @@ -2513,7 +2616,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet return Ok(None); } - let workspace_path = match self.require_main_session_workspace(session_id) { + let storage_path = match self.require_main_session_storage_path(session_id).await { Ok(path) => path, Err(_) => return Ok(None), }; @@ -2523,14 +2626,15 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .turn_cumulative_billable_tokens(source_turn_id); let goal_before = self - .get_thread_goal(session_id, workspace_path.as_path()) + .thread_goal_store() + .get_thread_goal(session_id, storage_path.as_path()) .await?; let plan = maybe_build_continuation_after_turn( &self.thread_goal_store(), self.thread_goal_runtime.as_ref(), session_id, - workspace_path.as_path(), + storage_path.as_path(), source_turn_id, turn_tokens, turn_completed, @@ -2538,7 +2642,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .await?; let goal_after = self - .get_thread_goal(session_id, workspace_path.as_path()) + .thread_goal_store() + .get_thread_goal(session_id, storage_path.as_path()) .await?; if goal_before.as_ref().map(|goal| goal.status) != goal_after.as_ref().map(|goal| goal.status) @@ -2588,14 +2693,9 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet }; if needs_restore { - let workspace_path = session.config.workspace_path.as_deref().ok_or_else(|| { - BitFunError::Validation(format!( - "workspace_path is required when restoring session: {}", - session_id - )) - })?; + let restore_path = self.restore_path_for_existing_session(&session_id).await?; self.session_manager - .restore_session(Path::new(workspace_path), &session_id) + .restore_session(&restore_path, &session_id) .await?; } @@ -2918,24 +3018,34 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet "Starting session history restore: session_id={}", session_id ); + let restore_workspace_path = session + .config + .workspace_path + .as_deref() + .or(workspace_path.as_deref()) + .ok_or_else(|| { + BitFunError::Validation(format!( + "workspace_path is required when restoring session: {}", + session_id + )) + })?; + let restore_path = Self::resolve_session_restore_path( + restore_workspace_path, + session + .config + .remote_connection_id + .as_deref() + .or(remote_connection_id.as_deref()), + session + .config + .remote_ssh_host + .as_deref() + .or(remote_ssh_host.as_deref()), + ) + .await?; match self .session_manager - .restore_session( - Path::new( - session - .config - .workspace_path - .as_deref() - .or(workspace_path.as_deref()) - .ok_or_else(|| { - BitFunError::Validation(format!( - "workspace_path is required when restoring session: {}", - session_id - )) - })?, - ), - &session_id, - ) + .restore_session(&restore_path, &session_id) .await { Ok(_) => { @@ -3230,7 +3340,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet // remote_ssh_host (which would silently fall back to a slugified raw remote path). let session_storage_path = session_workspace .as_ref() - .map(|workspace| workspace.session_storage_path().to_path_buf()); + .map(|workspace| workspace.session_storage_dir().to_path_buf()); let runtime_tool_restrictions = if is_miniapp_headless_agent_run( user_message_metadata.as_ref(), @@ -3904,15 +4014,6 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet self.session_manager.list_sessions(workspace_path).await } - pub async fn resolve_session_workspace_path( - &self, - session_id: &str, - ) -> Option { - self.session_manager - .resolve_session_workspace_path(session_id) - .await - } - /// Get a best-effort message view for a session. pub async fn get_messages(&self, session_id: &str) -> BitFunResult> { self.session_manager.get_messages(session_id).await @@ -4478,7 +4579,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .map(|workspace| workspace.root_path_string()); let subagent_session_storage_path = subagent_workspace .as_ref() - .map(|workspace| workspace.session_storage_path().to_path_buf()); + .map(|workspace| workspace.session_storage_dir().to_path_buf()); let subagent_services = Self::build_workspace_services(&subagent_workspace).await; let execution_context = ExecutionContext { session_id: session_id.clone(), @@ -5342,6 +5443,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet })?; let parent_agent_type = parent_session.agent_type.clone(); let parent_workspace_path = parent_session.config.workspace_path.clone(); + let parent_remote_connection_id = parent_session.config.remote_connection_id.clone(); + let parent_remote_ssh_host = parent_session.config.remote_ssh_host.clone(); let background_task_id = format!("bg-subagent-{}", uuid::Uuid::new_v4()); let background_task_id_for_delivery = background_task_id.clone(); let task_description = request.user_input_text.clone(); @@ -5420,6 +5523,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet session_id: subagent_parent_info.session_id.clone(), agent_type: parent_agent_type, workspace_path: parent_workspace_path, + remote_connection_id: parent_remote_connection_id, + remote_ssh_host: parent_remote_ssh_host, content: delivery_text, display_content: Some(display_text), metadata, @@ -5788,7 +5893,7 @@ impl bitfun_runtime_ports::AgentSubmissionPort for ConversationCoordinator { return Ok(None); }; - self.restore_session(&binding.session_storage_path(), session_id) + self.restore_session(&binding.session_storage_dir(), session_id) .await .map(|session| Some(session.agent_type)) .map_err(|error| { @@ -5890,16 +5995,6 @@ impl bitfun_runtime_ports::AgentSessionManagementPort for ConversationCoordinato }) } - async fn resolve_session_workspace_path( - &self, - request: bitfun_runtime_ports::AgentSessionWorkspaceRequest, - ) -> bitfun_runtime_ports::PortResult> { - Ok(self - .resolve_session_workspace_path(&request.session_id) - .await - .map(|path| path.to_string_lossy().into_owned())) - } - async fn resolve_session_workspace_binding( &self, request: bitfun_runtime_ports::AgentSessionWorkspaceRequest, @@ -6422,6 +6517,14 @@ mod tests { let workspace_path = std::env::temp_dir().join(format!("bitfun-btw-baseline-test-{}", uuid::Uuid::new_v4())); std::fs::create_dir_all(&workspace_path).expect("workspace dir should exist"); + struct TempWorkspaceGuard(std::path::PathBuf); + impl Drop for TempWorkspaceGuard { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + let _workspace_guard = TempWorkspaceGuard(workspace_path.clone()); + let parent_session = session_manager .create_session( "Parent".to_string(), diff --git a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs index 1517eaf04..c0aac883f 100644 --- a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs +++ b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs @@ -205,6 +205,8 @@ impl DialogScheduler { session_id: String, agent_type: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, goal: ThreadGoal, ) -> Result<(), String> { let plan = build_thread_goal_resumed_delivery_plan(&goal); @@ -241,8 +243,8 @@ impl DialogScheduler { None, agent_type, workspace_path, - None, - None, + remote_connection_id, + remote_ssh_host, DialogSubmissionPolicy::new( DialogTriggerSource::AgentSession, queue_priority, @@ -265,6 +267,8 @@ impl DialogScheduler { session_id: String, agent_type: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, goal: ThreadGoal, ) -> Result<(), String> { let plan = build_thread_goal_objective_updated_delivery_plan(&goal); @@ -301,8 +305,8 @@ impl DialogScheduler { None, agent_type, workspace_path, - None, - None, + remote_connection_id, + remote_ssh_host, DialogSubmissionPolicy::new( DialogTriggerSource::AgentSession, queue_priority, @@ -329,6 +333,8 @@ impl DialogScheduler { session_id: String, agent_type: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, content: String, display_content: Option, user_message_metadata: Option, @@ -366,8 +372,8 @@ impl DialogScheduler { None, agent_type, workspace_path, - None, - None, + remote_connection_id, + remote_ssh_host, DialogSubmissionPolicy::new( DialogTriggerSource::AgentSession, queue_priority, @@ -386,10 +392,17 @@ impl DialogScheduler { &self, session_id: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, policy: DialogSubmissionPolicy, ) -> Result { let agent_type = self - .resolve_session_agent_type(&session_id, workspace_path.as_deref(), None, None) + .resolve_session_agent_type( + &session_id, + workspace_path.as_deref(), + remote_connection_id.as_deref(), + remote_ssh_host.as_deref(), + ) .await?; let (user_input, prepended_messages) = build_init_agents_md_user_input() .await @@ -402,8 +415,8 @@ impl DialogScheduler { None, agent_type, workspace_path, - None, - None, + remote_connection_id, + remote_ssh_host, policy, None, None, @@ -876,6 +889,8 @@ impl DialogScheduler { ActiveDialogTurn::new( resolved.clone(), queued_turn.workspace_path.clone(), + queued_turn.remote_connection_id.clone(), + queued_turn.remote_ssh_host.clone(), queued_turn.agent_type.clone(), queued_turn .original_user_input @@ -1043,8 +1058,8 @@ impl DialogScheduler { None, active_turn.agent_type_owned(), active_turn.workspace_path_owned(), - None, - None, + active_turn.remote_connection_id_owned(), + active_turn.remote_ssh_host_owned(), DialogSubmissionPolicy::for_source( DialogTriggerSource::AgentSession, ), @@ -1295,6 +1310,8 @@ impl AgentLifecycleDeliveryPort for DialogScheduler { request.session_id, request.agent_type, request.workspace_path, + request.remote_connection_id, + request.remote_ssh_host, request.content, request.display_content, metadata, @@ -1311,6 +1328,8 @@ impl AgentLifecycleDeliveryPort for DialogScheduler { request.session_id, request.agent_type, request.workspace_path, + request.remote_connection_id, + request.remote_ssh_host, request.goal, ) .await @@ -1321,6 +1340,8 @@ impl AgentLifecycleDeliveryPort for DialogScheduler { request.session_id, request.agent_type, request.workspace_path, + request.remote_connection_id, + request.remote_ssh_host, request.goal, ) .await @@ -1418,6 +1439,8 @@ mod tests { ActiveDialogTurn::new( "turn_1".to_string(), Some("/workspace".to_string()), + None, + None, "agentic".to_string(), "hello".to_string(), None, diff --git a/src/crates/assembly/core/src/agentic/deep_review/report.rs b/src/crates/assembly/core/src/agentic/deep_review/report.rs index 3acb40e82..5322d3395 100644 --- a/src/crates/assembly/core/src/agentic/deep_review/report.rs +++ b/src/crates/assembly/core/src/agentic/deep_review/report.rs @@ -149,10 +149,10 @@ pub(crate) async fn persist_deep_review_cache( let Some(coordinator) = get_global_coordinator() else { return Ok(()); }; - let session_storage_path = workspace.session_storage_path(); + let session_storage_dir = workspace.session_storage_dir(); let session_manager = coordinator.get_session_manager(); let Some(mut metadata) = session_manager - .load_session_metadata(&session_storage_path, session_id) + .load_session_metadata(&session_storage_dir, session_id) .await? else { return Ok(()); @@ -160,6 +160,6 @@ pub(crate) async fn persist_deep_review_cache( set_deep_review_cache(&mut metadata, cache_value); session_manager - .save_session_metadata(&session_storage_path, &metadata) + .save_session_metadata(&session_storage_dir, &metadata) .await } diff --git a/src/crates/assembly/core/src/agentic/insights/session_paths.rs b/src/crates/assembly/core/src/agentic/insights/session_paths.rs index e7ad35323..10824a89e 100644 --- a/src/crates/assembly/core/src/agentic/insights/session_paths.rs +++ b/src/crates/assembly/core/src/agentic/insights/session_paths.rs @@ -1,22 +1,13 @@ //! Resolve on-disk session roots for insights (local + remote SSH mirror). -use crate::infrastructure::get_path_manager_arc; -use crate::service::remote_ssh::workspace_state::get_effective_session_path; +use crate::agentic::session::session_store_port::CoreSessionStorePort; use crate::service::workspace::{get_global_workspace_service, WorkspaceInfo}; +use bitfun_runtime_ports::{SessionStoragePathRequest, SessionStorePort}; use std::collections::HashSet; use std::path::PathBuf; -/// Resolve the workspace path to pass to [`PersistenceManager`] for session lookups. -/// -/// For local workspaces this is the workspace root path itself — the persistence layer -/// derives the actual sessions directory via [`PathManager::project_sessions_dir`]. -/// For remote workspaces this is the local SSH mirror directory, which the persistence -/// layer treats as the storage root directly. -pub async fn effective_session_storage_path_for_workspace(ws: &WorkspaceInfo) -> PathBuf { - if ws.remote_ssh_connection_id().is_none() { - return ws.root_path.clone(); - } - +/// Resolve the final sessions directory for a tracked workspace. +pub async fn effective_session_storage_dir_for_workspace(ws: &WorkspaceInfo) -> PathBuf { let path_str = ws.root_path.to_string_lossy().to_string(); let conn = ws.remote_ssh_connection_id().map(|s| s.to_string()); let mut host = ws @@ -34,7 +25,15 @@ pub async fn effective_session_storage_path_for_workspace(ws: &WorkspaceInfo) -> } } - get_effective_session_path(&path_str, conn.as_deref(), host.as_deref()).await + CoreSessionStorePort::default() + .resolve_session_storage_path(SessionStoragePathRequest { + workspace_path: ws.root_path.clone(), + remote_connection_id: conn, + remote_ssh_host: host, + }) + .await + .map(|resolution| resolution.effective_storage_path) + .unwrap_or_else(|_| PathBuf::from(path_str)) } /// Unique workspace paths whose persisted session directories exist on disk. @@ -48,22 +47,11 @@ pub async fn collect_effective_session_storage_roots() -> Vec { return paths; }; - let path_manager = get_path_manager_arc(); - for ws in ws_service.list_workspace_infos().await { - let workspace_path = effective_session_storage_path_for_workspace(&ws).await; - - // For local workspaces the actual sessions directory is derived from the - // workspace root via the path manager. For remote workspaces the mirror - // directory itself is the sessions root. - let sessions_dir = if ws.remote_ssh_connection_id().is_none() { - path_manager.project_sessions_dir(&workspace_path) - } else { - workspace_path.clone() - }; + let sessions_dir = effective_session_storage_dir_for_workspace(&ws).await; - if sessions_dir.exists() && seen.insert(workspace_path.clone()) { - paths.push(workspace_path); + if sessions_dir.exists() && seen.insert(sessions_dir.clone()) { + paths.push(sessions_dir); } } diff --git a/src/crates/assembly/core/src/agentic/persistence/manager.rs b/src/crates/assembly/core/src/agentic/persistence/manager.rs index 968491780..33980852d 100644 --- a/src/crates/assembly/core/src/agentic/persistence/manager.rs +++ b/src/crates/assembly/core/src/agentic/persistence/manager.rs @@ -303,27 +303,33 @@ impl PersistenceManager { /// Resolve the on-disk sessions directory for `workspace_path`. /// - /// For local workspaces this delegates to `PathManager::project_sessions_dir`, - /// which slugifies the workspace root under `~/.bitfun/projects/`. - /// - /// For remote SSH workspaces, callers (notably `desktop_effective_session_storage_path`) - /// pass an already-resolved mirror path under `~/.bitfun/remote_ssh/{host}/{path}/sessions`. - /// In that case we MUST use the path as-is; otherwise the slug pipeline would treat the - /// mirror path as a workspace root and write/read to a bogus - /// `~/.bitfun/projects//sessions/` location. + /// Callers may pass either a logical workspace root or an already-resolved + /// managed sessions directory. Local workspace roots are slugified under + /// `~/.bitfun/projects/`; already-resolved local/remote sessions + /// directories are used as-is. fn project_sessions_dir(&self, workspace_path: &Path) -> PathBuf { - let remote_mirror_root = PathManager::remote_ssh_mirror_root(); - if workspace_path.starts_with(&remote_mirror_root) { - // Already resolved: either the mirror runtime root, the mirror sessions dir, - // or a session sub-dir. Treat the path as the sessions root directly. - // (Inputs that already include a trailing `sessions` segment stay correct; - // inputs at the mirror runtime root would historically fall back to the - // legacy slug, but no current call-site uses that shape.) + if self.is_resolved_sessions_dir(workspace_path) { return workspace_path.to_path_buf(); } self.path_manager.project_sessions_dir(workspace_path) } + fn is_resolved_sessions_dir(&self, path: &Path) -> bool { + if path.file_name().and_then(|value| value.to_str()) != Some("sessions") { + return false; + } + + let remote_mirror_root = self.path_manager.remote_ssh_mirror_root_dir(); + if path.starts_with(&remote_mirror_root) { + return true; + } + + let projects_root = self.path_manager.projects_root(); + path.parent() + .and_then(|runtime_root| runtime_root.parent()) + .is_some_and(|candidate| candidate == projects_root.as_path()) + } + fn metadata_path(&self, workspace_path: &Path, session_id: &str) -> PathBuf { self.session_layout(workspace_path) .metadata_path(session_id) @@ -410,8 +416,7 @@ impl PersistenceManager { } async fn ensure_runtime_for_write(&self, workspace_path: &Path) -> BitFunResult<()> { - let remote_mirror_root = PathManager::remote_ssh_mirror_root(); - if workspace_path.starts_with(&remote_mirror_root) { + if self.is_resolved_sessions_dir(workspace_path) { return Ok(()); } @@ -2949,8 +2954,8 @@ mod tests { #[tokio::test] async fn load_session_with_turns_returns_session_and_persisted_turns() { let workspace = TestWorkspace::new(); - let manager = PersistenceManager::new(Arc::new(PathManager::new().expect("path manager"))) - .expect("persistence manager"); + let manager = + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"); let session_id = Uuid::new_v4().to_string(); let session = Session::new_with_id( session_id.clone(), @@ -3046,8 +3051,8 @@ mod tests { #[tokio::test] async fn save_dialog_turn_updates_metadata_without_scanning_unrelated_turn_files() { let workspace = TestWorkspace::new(); - let manager = PersistenceManager::new(Arc::new(PathManager::new().expect("path manager"))) - .expect("persistence manager"); + let manager = + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"); let session_id = Uuid::new_v4().to_string(); let session = Session::new_with_id( session_id.clone(), @@ -3122,8 +3127,8 @@ mod tests { #[tokio::test] async fn concurrent_dialog_turn_saves_keep_metadata_counts_consistent() { let workspace = TestWorkspace::new(); - let manager = PersistenceManager::new(Arc::new(PathManager::new().expect("path manager"))) - .expect("persistence manager"); + let manager = + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"); let session_id = Uuid::new_v4().to_string(); let session = Session::new_with_id( session_id.clone(), @@ -3550,11 +3555,61 @@ mod tests { assert!(runtime.layout_state_file.exists()); } + #[tokio::test] + async fn local_sessions_dir_input_is_used_without_reslugging() { + let workspace = TestWorkspace::new(); + let path_manager = workspace.path_manager(); + let sessions_dir = path_manager.project_sessions_dir(workspace.path()); + let manager = PersistenceManager::new(path_manager).expect("persistence manager"); + + let metadata = SessionMetadata::new( + Uuid::new_v4().to_string(), + "Resolved sessions root".to_string(), + "agent".to_string(), + "model".to_string(), + ); + + manager + .save_session_metadata(&sessions_dir, &metadata) + .await + .expect("metadata should save under resolved sessions dir"); + + assert_eq!( + manager.index_path(&sessions_dir), + sessions_dir.join("index.json") + ); + assert!(sessions_dir + .join(&metadata.session_id) + .join("metadata.json") + .exists()); + } + + #[tokio::test] + async fn remote_sessions_dir_input_is_used_without_accepting_runtime_root() { + let test_root = + std::env::temp_dir().join(format!("bitfun-persistence-test-{}", Uuid::new_v4())); + let path_manager = Arc::new(PathManager::with_user_root_for_tests( + test_root.join("user"), + )); + let manager = PersistenceManager::new(path_manager.clone()).expect("persistence manager"); + let runtime_root = path_manager + .remote_ssh_mirror_root_dir() + .join("example-host") + .join("root") + .join("repo"); + let sessions_dir = runtime_root.join("sessions"); + + assert_eq!(manager.project_sessions_dir(&sessions_dir), sessions_dir); + assert_ne!(manager.project_sessions_dir(&runtime_root), runtime_root); + + let _ = std::fs::remove_dir_all(&test_root); + } + #[tokio::test] async fn skill_agent_snapshots_persist_and_truncate_with_context_snapshots() { let workspace = TestWorkspace::new(); - let manager = PersistenceManager::new(Arc::new(PathManager::new().expect("path manager"))) - .expect("persistence manager"); + let manager = + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"); let session_id = Uuid::new_v4().to_string(); let snapshot = TurnSkillAgentSnapshot { skills: vec![SkillSnapshotEntry { diff --git a/src/crates/assembly/core/src/agentic/session/session_manager.rs b/src/crates/assembly/core/src/agentic/session/session_manager.rs index a6e973324..4fdafb7bd 100644 --- a/src/crates/assembly/core/src/agentic/session/session_manager.rs +++ b/src/crates/assembly/core/src/agentic/session/session_manager.rs @@ -415,11 +415,53 @@ impl SessionManager { .unwrap_or(true) } - /// Resolve the effective storage path for a session's workspace. - async fn effective_workspace_path_from_config(config: &SessionConfig) -> Option { - CoreSessionStorePort::resolve_storage_path_for_config(config) + async fn effective_storage_path_for_config_with_persistence( + persistence_manager: &PersistenceManager, + config: &SessionConfig, + ) -> Option { + let workspace_path = config.workspace_path.as_ref()?; + let identity = + crate::service::remote_ssh::workspace_state::resolve_workspace_session_identity( + workspace_path, + config.remote_connection_id.as_deref(), + config.remote_ssh_host.as_deref(), + ) + .await?; + + let runtime_service = persistence_manager.runtime_service(); + Some(if identity.hostname == LOCAL_WORKSPACE_SSH_HOST { + runtime_service + .context_for_local_workspace(Path::new(identity.logical_workspace_path())) + .sessions_dir + } else if identity.hostname == "_unresolved" { + bitfun_services_integrations::remote_ssh::unresolved_remote_session_storage_dir( + runtime_service.path_manager().remote_ssh_mirror_root_dir(), + identity.remote_connection_id.as_deref().unwrap_or_default(), + identity.logical_workspace_path(), + ) + } else { + runtime_service + .context_for_remote_workspace(&identity.hostname, identity.logical_workspace_path()) + .sessions_dir + }) + } + + async fn effective_storage_path_for_config(&self, config: &SessionConfig) -> Option { + Self::effective_storage_path_for_config_with_persistence( + self.persistence_manager.as_ref(), + config, + ) + .await + } + + async fn effective_storage_path_for_workspace_path(&self, workspace_path: &Path) -> PathBuf { + let tmp_config = SessionConfig { + workspace_path: Some(workspace_path.to_string_lossy().to_string()), + ..Default::default() + }; + self.effective_storage_path_for_config(&tmp_config) .await - .map(|resolution| resolution.effective_storage_path) + .unwrap_or_else(|| workspace_path.to_path_buf()) } #[allow(dead_code)] @@ -433,58 +475,7 @@ impl SessionManager { /// For remote workspaces, maps the remote path to a local session storage path. async fn effective_session_workspace_path(&self, session_id: &str) -> Option { let config = self.sessions.get(session_id)?.config.clone(); - Self::effective_workspace_path_from_config(&config).await - } - - /// Resolve the logical workspace path bound to a session. - /// - /// This prefers the in-memory session config, then the persisted metadata - /// reachable via the session workspace index, and finally scans tracked - /// workspaces known to the global workspace service ordered by recent access. - pub async fn resolve_session_workspace_path(&self, session_id: &str) -> Option { - if let Some(workspace_path) = self - .get_session(session_id) - .and_then(|session| session.config.workspace_path) - .filter(|path| !path.is_empty()) - { - return Some(PathBuf::from(workspace_path)); - } - - let indexed_workspace_path = self - .session_workspace_index - .get(session_id) - .map(|entry| entry.clone()); - if let Some(workspace_path) = indexed_workspace_path { - match self - .persistence_manager - .load_session_metadata(&workspace_path, session_id) - .await - { - Ok(Some(metadata)) => { - if let Some(bound_workspace) = - metadata.workspace_path.filter(|path| !path.is_empty()) - { - return Some(PathBuf::from(bound_workspace)); - } - return Some(workspace_path); - } - Ok(None) => {} - Err(err) => { - debug!( - "Failed to load indexed session metadata while resolving workspace: session_id={} workspace={} error={}", - session_id, - workspace_path.display(), - err - ); - } - } - } - - if let Some(binding) = self.resolve_session_workspace_binding(session_id).await { - return Some(binding.root_path().to_path_buf()); - } - - None + self.effective_storage_path_for_config(&config).await } pub async fn resolve_session_workspace_binding( @@ -1322,7 +1313,8 @@ impl SessionManager { BitFunError::Validation("Session workspace_path is required".to_string()) })?; - let session_storage_path = Self::effective_workspace_path_from_config(&config) + let session_storage_path = self + .effective_storage_path_for_config(&config) .await .ok_or_else(|| { BitFunError::Validation("Session workspace_path is required".to_string()) @@ -2686,15 +2678,9 @@ impl SessionManager { }; let restore_started_at = Instant::now(); let storage_path_started_at = Instant::now(); - let session_storage_path = CoreSessionStorePort - .resolve_session_storage_path(SessionStoragePathRequest { - workspace_path: restore_request.workspace_path.clone(), - remote_connection_id: None, - remote_ssh_host: None, - }) - .await - .map(|resolution| resolution.effective_storage_path) - .unwrap_or_else(|_| restore_request.workspace_path.clone()); + let session_storage_path = self + .effective_storage_path_for_workspace_path(&restore_request.workspace_path) + .await; let resolve_storage_path_duration_ms = elapsed_ms_u64(storage_path_started_at); debug!( "Session view restore phase completed: session_id={}, phase=resolve_storage_path, duration_ms={}", @@ -2830,16 +2816,9 @@ impl SessionManager { let session_already_in_memory = self.sessions.contains_key(session_id); let storage_path_started_at = Instant::now(); - let session_storage_path = { - let ws = workspace_path.to_string_lossy().to_string(); - let tmp_config = SessionConfig { - workspace_path: Some(ws), - ..Default::default() - }; - Self::effective_workspace_path_from_config(&tmp_config) - .await - .unwrap_or_else(|| workspace_path.to_path_buf()) - }; + let session_storage_path = self + .effective_storage_path_for_workspace_path(workspace_path) + .await; debug!( "Session restore phase completed: session_id={}, phase=resolve_storage_path, duration_ms={}", session_id, @@ -3449,7 +3428,8 @@ impl SessionManager { let session = self .get_session(session_id) .ok_or_else(|| BitFunError::NotFound(format!("Session not found: {}", session_id)))?; - let workspace_path = Self::effective_workspace_path_from_config(&session.config) + let workspace_path = self + .effective_storage_path_for_config(&session.config) .await .ok_or_else(|| { BitFunError::Validation(format!( @@ -3678,7 +3658,8 @@ impl SessionManager { let session = self .get_session(session_id) .ok_or_else(|| BitFunError::NotFound(format!("Session not found: {}", session_id)))?; - let workspace_path = Self::effective_workspace_path_from_config(&session.config) + let workspace_path = self + .effective_storage_path_for_config(&session.config) .await .ok_or_else(|| { BitFunError::Validation(format!( @@ -4409,7 +4390,11 @@ impl SessionManager { continue; } if let Some(workspace_path) = - Self::effective_workspace_path_from_config(&snapshot.session.config).await + Self::effective_storage_path_for_config_with_persistence( + persistence.as_ref(), + &snapshot.session.config, + ) + .await { if !Self::auto_save_snapshot_is_current(&sessions, &snapshot) { continue; @@ -4471,7 +4456,11 @@ impl SessionManager { if enable_persistence && Self::should_persist_session(&session) { if let Some(workspace_path) = - Self::effective_workspace_path_from_config(&session.config).await + Self::effective_storage_path_for_config_with_persistence( + persistence.as_ref(), + &session.config, + ) + .await { if Self::cleanup_snapshot_for_candidate( &sessions, @@ -4528,7 +4517,6 @@ mod tests { use crate::service::config::types::{ AIConfig as ServiceAIConfig, AIModelConfig as ServiceAIModelConfig, }; - use crate::service::remote_ssh::workspace_state::local_workspace_roots_equal; use crate::service::session::{ DialogTurnData, DialogTurnKind, ModelRoundData, SessionKind, SessionMetadata, SessionRelationship, SessionRelationshipKind, ToolCallData, ToolItemData, ToolResultData, @@ -4596,11 +4584,17 @@ mod tests { ) } + fn test_path_manager() -> Arc { + let root = + std::env::temp_dir().join(format!("bitfun-session-manager-test-{}", Uuid::new_v4())); + Arc::new(PathManager::with_user_root_for_tests( + root.join("user-root"), + )) + } + fn in_memory_test_manager() -> SessionManager { - let persistence_manager = Arc::new( - PersistenceManager::new(Arc::new(PathManager::new().expect("path manager"))) - .expect("persistence manager"), - ); + let persistence_manager = + Arc::new(PersistenceManager::new(test_path_manager()).expect("persistence manager")); SessionManager::new( Arc::new(SessionContextStore::new()), persistence_manager, @@ -5174,13 +5168,53 @@ mod tests { ); } + #[tokio::test] + async fn core_session_store_port_resolves_local_storage_to_sessions_dir() { + use bitfun_runtime_ports::{ + SessionStorageKind, SessionStoragePathRequest, SessionStorePort, + }; + + let workspace = TestWorkspace::new(); + let path_manager = workspace.path_manager(); + let port = CoreSessionStorePort::with_path_manager_for_tests(path_manager.clone()); + let resolution = port + .resolve_session_storage_path(SessionStoragePathRequest { + workspace_path: workspace.path().to_path_buf(), + remote_connection_id: None, + remote_ssh_host: None, + }) + .await + .expect("storage path should resolve"); + + assert_eq!(resolution.storage_kind, SessionStorageKind::Local); + assert_eq!( + resolution.effective_storage_path, + path_manager.project_sessions_dir(workspace.path()) + ); + assert_ne!(resolution.effective_storage_path, workspace.path()); + + let resolved_again = port + .resolve_session_storage_path(SessionStoragePathRequest { + workspace_path: resolution.effective_storage_path.clone(), + remote_connection_id: None, + remote_ssh_host: None, + }) + .await + .expect("resolved sessions dir should pass through"); + assert_eq!( + resolved_again.effective_storage_path, + resolution.effective_storage_path + ); + } + #[tokio::test] async fn core_session_store_port_resolves_unresolved_remote_storage_path() { use bitfun_runtime_ports::{ SessionStorageKind, SessionStoragePathRequest, SessionStorePort, }; - let port = CoreSessionStorePort; + let workspace = TestWorkspace::new(); + let port = CoreSessionStorePort::with_path_manager_for_tests(workspace.path_manager()); let resolution = port .resolve_session_storage_path(SessionStoragePathRequest { workspace_path: PathBuf::from("/remote/project"), @@ -5202,12 +5236,58 @@ mod tests { ); } + #[tokio::test] + async fn core_session_store_port_resolved_remote_sessions_dir_passes_through_only_sessions_root( + ) { + use bitfun_runtime_ports::{ + SessionStorageKind, SessionStoragePathRequest, SessionStorePort, + }; + + let workspace = TestWorkspace::new(); + let path_manager = workspace.path_manager(); + let port = CoreSessionStorePort::with_path_manager_for_tests(path_manager.clone()); + let sessions_dir = + bitfun_services_integrations::remote_ssh::remote_workspace_session_mirror_dir( + path_manager.remote_ssh_mirror_root_dir(), + "example-host", + "/root/repo", + ); + let resolved = port + .resolve_session_storage_path(SessionStoragePathRequest { + workspace_path: sessions_dir.clone(), + remote_connection_id: None, + remote_ssh_host: None, + }) + .await + .expect("resolved remote sessions dir should pass through"); + + assert_eq!(resolved.storage_kind, SessionStorageKind::Remote); + assert_eq!(resolved.effective_storage_path, sessions_dir); + + let runtime_root = bitfun_services_integrations::remote_ssh::remote_workspace_runtime_root( + path_manager.remote_ssh_mirror_root_dir(), + "example-host", + "/root/repo", + ); + let runtime_root_resolution = port + .resolve_session_storage_path(SessionStoragePathRequest { + workspace_path: runtime_root.clone(), + remote_connection_id: None, + remote_ssh_host: None, + }) + .await; + + assert!( + runtime_root_resolution.is_err(), + "remote runtime root must not pass as a resolved sessions dir" + ); + } + #[tokio::test] async fn restore_session_view_loads_turns_without_restoring_runtime_context() { let workspace = TestWorkspace::new(); let persistence_manager = Arc::new( - PersistenceManager::new(Arc::new(PathManager::new().expect("path manager"))) - .expect("persistence manager"), + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), ); let manager = test_manager(persistence_manager.clone()); let session_id = Uuid::new_v4().to_string(); @@ -5347,8 +5427,7 @@ mod tests { async fn restore_session_view_preserves_full_visible_tool_result_payload() { let workspace = TestWorkspace::new(); let persistence_manager = Arc::new( - PersistenceManager::new(Arc::new(PathManager::new().expect("path manager"))) - .expect("persistence manager"), + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), ); let manager = test_manager(persistence_manager.clone()); let session_id = Uuid::new_v4().to_string(); @@ -6089,7 +6168,13 @@ mod tests { #[tokio::test] async fn delete_session_removes_workspace_cache_entry() { let workspace = TestWorkspace::new(); - let manager = in_memory_test_manager(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let expected_storage_path = persistence_manager + .path_manager() + .project_sessions_dir(workspace.path()); + let manager = test_manager(persistence_manager); let session = manager .create_session( "Cached session".to_string(), @@ -6107,8 +6192,8 @@ mod tests { .session_workspace_index .get(&session.session_id) .as_deref() - .map(|entry| local_workspace_roots_equal(entry, workspace.path())), - Some(true) + .map(|entry| entry.to_path_buf()), + Some(expected_storage_path) ); manager @@ -6204,10 +6289,8 @@ mod tests { #[tokio::test] async fn records_subagent_partial_timeout_in_evidence_ledger() { - let persistence_manager = Arc::new( - PersistenceManager::new(Arc::new(PathManager::new().expect("path manager"))) - .expect("persistence manager"), - ); + let persistence_manager = + Arc::new(PersistenceManager::new(test_path_manager()).expect("persistence manager")); let manager = test_manager(persistence_manager); let event = manager.record_subagent_partial_timeout( diff --git a/src/crates/assembly/core/src/agentic/session/session_store_port.rs b/src/crates/assembly/core/src/agentic/session/session_store_port.rs index 283154ddd..607994b6c 100644 --- a/src/crates/assembly/core/src/agentic/session/session_store_port.rs +++ b/src/crates/assembly/core/src/agentic/session/session_store_port.rs @@ -1,4 +1,5 @@ -use std::path::PathBuf; +use std::path::{Path, PathBuf}; +use std::sync::Arc; use bitfun_runtime_ports::{ PortError, PortErrorKind, PortResult, RuntimeServiceCapability, RuntimeServicePort, @@ -6,15 +7,32 @@ use bitfun_runtime_ports::{ }; use crate::agentic::core::SessionConfig; +use crate::infrastructure::{get_path_manager_arc, PathManager}; use crate::service::remote_ssh::workspace_state::{ resolve_workspace_session_identity, unresolved_remote_session_storage_dir, LOCAL_WORKSPACE_SSH_HOST, }; +use crate::service::WorkspaceRuntimeService; -#[derive(Debug, Clone, Copy, Default)] -pub struct CoreSessionStorePort; +#[derive(Debug, Clone, Default)] +pub struct CoreSessionStorePort { + path_manager: Option>, +} impl CoreSessionStorePort { + #[cfg(test)] + pub fn with_path_manager_for_tests(path_manager: Arc) -> Self { + Self { + path_manager: Some(path_manager), + } + } + + fn path_manager(&self) -> Arc { + self.path_manager + .clone() + .unwrap_or_else(get_path_manager_arc) + } + pub async fn resolve_storage_path_for_config( config: &SessionConfig, ) -> Option { @@ -29,6 +47,35 @@ impl CoreSessionStorePort { .await .ok() } + + fn resolved_sessions_dir_kind( + path_manager: &PathManager, + path: &std::path::Path, + ) -> Option { + if path.file_name().and_then(|value| value.to_str()) != Some("sessions") { + return None; + } + + let remote_mirror_root = path_manager.remote_ssh_mirror_root_dir(); + if path.starts_with(&remote_mirror_root) { + return Some( + if path + .components() + .any(|component| component.as_os_str() == std::ffi::OsStr::new("_unresolved")) + { + SessionStorageKind::UnresolvedRemote + } else { + SessionStorageKind::Remote + }, + ); + } + + let projects_root = path_manager.projects_root(); + path.parent() + .and_then(|runtime_root| runtime_root.parent()) + .is_some_and(|candidate| candidate == projects_root.as_path()) + .then_some(SessionStorageKind::Local) + } } impl RuntimeServicePort for CoreSessionStorePort { @@ -43,6 +90,19 @@ impl SessionStorePort for CoreSessionStorePort { &self, request: SessionStoragePathRequest, ) -> PortResult { + let path_manager = self.path_manager(); + if let Some(storage_kind) = + Self::resolved_sessions_dir_kind(&path_manager, &request.workspace_path) + { + return Ok(SessionStoragePathResolution::new( + request.workspace_path.clone(), + request.workspace_path, + storage_kind, + request.remote_connection_id, + request.remote_ssh_host, + )); + } + let workspace_path = request.workspace_path.to_string_lossy().to_string(); let identity = resolve_workspace_session_identity( &workspace_path, @@ -58,10 +118,13 @@ impl SessionStorePort for CoreSessionStorePort { })?; let requested_workspace_path = request.workspace_path; + let runtime_service = WorkspaceRuntimeService::new(path_manager); let (effective_storage_path, storage_kind, remote_ssh_host) = if identity.hostname == LOCAL_WORKSPACE_SSH_HOST { ( - PathBuf::from(identity.logical_workspace_path()), + runtime_service + .context_for_local_workspace(Path::new(identity.logical_workspace_path())) + .sessions_dir, SessionStorageKind::Local, None, ) @@ -76,7 +139,12 @@ impl SessionStorePort for CoreSessionStorePort { ) } else { ( - identity.session_storage_path(), + runtime_service + .context_for_remote_workspace( + &identity.hostname, + identity.logical_workspace_path(), + ) + .sessions_dir, SessionStorageKind::Remote, Some(identity.hostname.clone()), ) diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/bash_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/bash_tool.rs index facd03e34..065384acb 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/bash_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/bash_tool.rs @@ -54,6 +54,8 @@ async fn deliver_background_bash_result( parent_session_id: String, parent_agent_type: String, parent_workspace_path: Option, + parent_remote_connection_id: Option, + parent_remote_ssh_host: Option, delivery_text: String, display_text: String, metadata: serde_json::Map, @@ -76,6 +78,8 @@ async fn deliver_background_bash_result( session_id: parent_session_id.clone(), agent_type: parent_agent_type, workspace_path: parent_workspace_path, + remote_connection_id: parent_remote_connection_id, + remote_ssh_host: parent_remote_ssh_host, content: delivery_text, display_content: Some(display_text), metadata, @@ -1107,6 +1111,16 @@ impl BashTool { let parent_workspace_path = context .workspace_root() .map(|path| path.to_string_lossy().to_string()); + let parent_remote_connection_id = context + .workspace + .as_ref() + .and_then(|workspace| workspace.connection_id().map(ToOwned::to_owned)); + let parent_remote_ssh_host = context + .workspace + .as_ref() + .filter(|workspace| workspace.is_remote()) + .map(|workspace| workspace.session_identity.hostname.clone()) + .filter(|value| !value.trim().is_empty()); let command = command_str.to_string(); let working_directory = initial_cwd.to_string(); let terminal_session_id = bg_session_id.clone(); @@ -1230,6 +1244,8 @@ impl BashTool { parent_session_id.clone(), parent_agent_type.clone(), parent_workspace_path.clone(), + parent_remote_connection_id.clone(), + parent_remote_ssh_host.clone(), delivery_text, display_text, json_object_metadata(metadata), @@ -1268,6 +1284,8 @@ impl BashTool { parent_session_id.clone(), parent_agent_type.clone(), parent_workspace_path.clone(), + parent_remote_connection_id.clone(), + parent_remote_ssh_host.clone(), delivery_text, display_text, json_object_metadata(metadata), @@ -1308,6 +1326,8 @@ impl BashTool { parent_session_id.clone(), parent_agent_type.clone(), parent_workspace_path.clone(), + parent_remote_connection_id.clone(), + parent_remote_ssh_host.clone(), delivery_text, display_text, json_object_metadata(metadata), diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/code_review_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/code_review_tool.rs index 0018db7ee..9094c7993 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/code_review_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/code_review_tool.rs @@ -631,10 +631,10 @@ impl Tool for CodeReviewTool { context.workspace.as_ref(), get_global_coordinator(), ) { - let session_storage_path = workspace.session_storage_path(); + let session_storage_dir = workspace.session_storage_dir(); match coordinator .get_session_manager() - .load_session_metadata(&session_storage_path, session_id) + .load_session_metadata(&session_storage_dir, session_id) .await { Ok(Some(metadata)) => { diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs index 6296f684b..f1c6a2110 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs @@ -261,11 +261,11 @@ Examples: )) })?; let display_workspace = workspace.root_path_string(); - let session_storage_path = workspace.session_storage_path(); + let session_storage_dir = workspace.session_storage_dir(); let manager = PersistenceManager::new(Arc::new(PathManager::new()?))?; let transcript = manager .export_session_transcript( - &session_storage_path, + &session_storage_dir, &session_id, &SessionTranscriptExportOptions { tools: params.tools.unwrap_or(false), diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/session_message_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/session_message_tool.rs index cd20ec679..cf882cd43 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/session_message_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/session_message_tool.rs @@ -201,6 +201,15 @@ impl SessionMessageTool { } } + fn same_workspace_identity( + left: &SessionMessageWorkspaceTarget, + right: &SessionMessageWorkspaceTarget, + ) -> bool { + left.workspace_path == right.workspace_path + && left.remote_connection_id == right.remote_connection_id + && left.remote_ssh_host == right.remote_ssh_host + } + fn format_forwarded_message( &self, message: &str, @@ -534,15 +543,7 @@ Allowed agent types when creating a session: let requested_workspace = self.resolve_workspace(workspace, context)?; let requested_target = self.workspace_target_from_context(requested_workspace.clone(), context); - let remote_mismatch = requested_target - .remote_connection_id - .as_deref() - .zip(workspace_target.remote_connection_id.as_deref()) - .map(|(left, right)| left != right) - .unwrap_or(false); - if requested_target.workspace_path != workspace_target.workspace_path - || remote_mismatch - { + if !Self::same_workspace_identity(&requested_target, &workspace_target) { return Err(BitFunError::NotFound(format!( "Session '{}' not found in workspace '{}'", target_session_id, requested_target.workspace_path @@ -723,6 +724,46 @@ mod tests { } } + fn workspace_target( + workspace_path: &str, + remote_connection_id: Option<&str>, + remote_ssh_host: Option<&str>, + ) -> SessionMessageWorkspaceTarget { + SessionMessageWorkspaceTarget { + workspace_path: workspace_path.to_string(), + remote_connection_id: remote_connection_id.map(ToOwned::to_owned), + remote_ssh_host: remote_ssh_host.map(ToOwned::to_owned), + } + } + + #[test] + fn workspace_identity_matches_full_remote_tuple() { + let left = workspace_target("/root/repo", Some("conn-1"), Some("host-a")); + let right = workspace_target("/root/repo", Some("conn-1"), Some("host-a")); + + assert!(SessionMessageTool::same_workspace_identity(&left, &right)); + } + + #[test] + fn workspace_identity_rejects_remote_local_parity_mismatch() { + let requested = workspace_target("/root/repo", None, None); + let target = workspace_target("/root/repo", Some("conn-1"), Some("host-a")); + + assert!(!SessionMessageTool::same_workspace_identity( + &requested, &target + )); + } + + #[test] + fn workspace_identity_rejects_remote_host_mismatch() { + let requested = workspace_target("/root/repo", Some("conn-1"), Some("host-a")); + let target = workspace_target("/root/repo", Some("conn-1"), Some("host-b")); + + assert!(!SessionMessageTool::same_workspace_identity( + &requested, &target + )); + } + #[tokio::test] async fn validate_existing_session_rejects_agent_type_override() { let tool = SessionMessageTool::new(); diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/task_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/task_tool.rs index b53e8ed8e..2f4fae3bc 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/task_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/task_tool.rs @@ -906,10 +906,10 @@ impl Tool for TaskTool { })?; let mut run_manifest = context.custom_data.get("deep_review_run_manifest").cloned(); if let Some(workspace) = context.workspace.as_ref() { - let session_storage_path = workspace.session_storage_path(); + let session_storage_dir = workspace.session_storage_dir(); match coordinator .get_session_manager() - .load_session_metadata(&session_storage_path, &session_id) + .load_session_metadata(&session_storage_dir, &session_id) .await { Ok(Some(metadata)) => { diff --git a/src/crates/assembly/core/src/agentic/tools/tool_context_runtime.rs b/src/crates/assembly/core/src/agentic/tools/tool_context_runtime.rs index f2035c0eb..4275759f7 100644 --- a/src/crates/assembly/core/src/agentic/tools/tool_context_runtime.rs +++ b/src/crates/assembly/core/src/agentic/tools/tool_context_runtime.rs @@ -492,6 +492,16 @@ impl ToolUseContext { } pub fn current_workspace_runtime_root(&self) -> BitFunResult { + #[cfg(test)] + if let Some(path) = self + .custom_data + .get("__bitfun_test_runtime_root") + .and_then(|value| value.as_str()) + .filter(|path| !path.trim().is_empty()) + { + return Ok(PathBuf::from(path)); + } + let workspace = self.workspace.as_ref().ok_or_else(|| { BitFunError::tool("A workspace is required to resolve runtime artifacts".to_string()) })?; diff --git a/src/crates/assembly/core/src/agentic/tools/tool_result_storage.rs b/src/crates/assembly/core/src/agentic/tools/tool_result_storage.rs index dc51850a3..468a11327 100644 --- a/src/crates/assembly/core/src/agentic/tools/tool_result_storage.rs +++ b/src/crates/assembly/core/src/agentic/tools/tool_result_storage.rs @@ -350,6 +350,12 @@ mod tests { use std::path::PathBuf; fn test_context(root: PathBuf) -> ToolUseContext { + let mut custom_data = HashMap::new(); + custom_data.insert( + "__bitfun_test_runtime_root".to_string(), + json!(root.join("runtime")), + ); + ToolUseContext { tool_call_id: Some("call_1".to_string()), agent_type: Some("agent".to_string()), @@ -357,7 +363,7 @@ mod tests { dialog_turn_id: Some("turn_1".to_string()), workspace: Some(WorkspaceBinding::new(None, root)), unlocked_collapsed_tools: Vec::new(), - custom_data: HashMap::new(), + custom_data, computer_use_host: None, runtime_tool_restrictions: ToolRuntimeRestrictions::default(), runtime_handles: bitfun_runtime_ports::ToolRuntimeHandles::default(), diff --git a/src/crates/assembly/core/src/agentic/workspace.rs b/src/crates/assembly/core/src/agentic/workspace.rs index 368805eba..1f9e5ad8a 100644 --- a/src/crates/assembly/core/src/agentic/workspace.rs +++ b/src/crates/assembly/core/src/agentic/workspace.rs @@ -1,4 +1,5 @@ use crate::service::remote_ssh::workspace_state::WorkspaceSessionIdentity; +use crate::service::workspace_runtime::WorkspaceRuntimeService; use async_trait::async_trait; pub use bitfun_runtime_ports::{ WorkspaceCommandOptions, WorkspaceCommandResult, WorkspaceDirEntry, WorkspaceFileSystem, @@ -79,6 +80,18 @@ impl WorkspaceBinding { self.root_path.to_string_lossy().to_string() } + /// Logical workspace root used by tools, display, and workspace-bound IO. + /// + /// For local workspaces this is the local project root. For remote SSH + /// workspaces this is the root path on the remote host. + pub fn logical_workspace_path(&self) -> &Path { + &self.root_path + } + + pub fn logical_workspace_path_string(&self) -> String { + self.logical_workspace_path().to_string_lossy().to_string() + } + pub fn is_remote(&self) -> bool { matches!(self.backend, WorkspaceBackend::Remote { .. }) } @@ -90,9 +103,30 @@ impl WorkspaceBinding { } } - /// The path to use for session persistence. - pub fn session_storage_path(&self) -> PathBuf { - self.session_identity.session_storage_path() + /// Final on-disk sessions directory for this workspace binding. + pub fn session_storage_dir(&self) -> PathBuf { + let runtime_service = + WorkspaceRuntimeService::new(crate::infrastructure::get_path_manager_arc()); + if self.is_remote() { + if self.session_identity.hostname == "_unresolved" { + if let Some(connection_id) = self.session_identity.remote_connection_id.as_deref() { + return crate::service::remote_ssh::workspace_state::unresolved_remote_session_storage_dir( + connection_id, + self.session_identity.logical_workspace_path(), + ); + } + } + return runtime_service + .context_for_remote_workspace( + &self.session_identity.hostname, + self.session_identity.logical_workspace_path(), + ) + .sessions_dir; + } + + runtime_service + .context_for_local_workspace(self.logical_workspace_path()) + .sessions_dir } } @@ -105,7 +139,7 @@ mod tests { use std::path::PathBuf; #[test] - fn remote_workspace_binding_uses_session_identity_storage_path() { + fn remote_workspace_binding_uses_session_identity_storage_dir() { let session_identity = workspace_session_identity( "/home/wsp/projects/test", Some("conn-1"), @@ -122,7 +156,7 @@ mod tests { assert!(matches!(binding.backend, WorkspaceBackend::Remote { .. })); assert_eq!( - binding.session_storage_path(), + binding.session_storage_dir(), remote_workspace_session_mirror_dir("127.0.0.1", "/home/wsp/projects/test") ); } diff --git a/src/crates/assembly/core/src/infrastructure/app_paths/path_manager.rs b/src/crates/assembly/core/src/infrastructure/app_paths/path_manager.rs index 45c7251f7..f2d271c51 100644 --- a/src/crates/assembly/core/src/infrastructure/app_paths/path_manager.rs +++ b/src/crates/assembly/core/src/infrastructure/app_paths/path_manager.rs @@ -276,9 +276,15 @@ impl PathManager { /// /// Session/chat persistence for SSH workspaces lives under /// `{this}/{sanitized_host}/{remote_path_segments}/sessions/`. + pub fn remote_ssh_mirror_root_dir(&self) -> PathBuf { + self.bitfun_home_dir().join("remote_ssh") + } + + /// Root for per-host, per-remote-path workspace mirrors using the default + /// process path manager. pub fn remote_ssh_mirror_root() -> PathBuf { Self::new() - .map(|pm| pm.bitfun_home_dir().join("remote_ssh")) + .map(|pm| pm.remote_ssh_mirror_root_dir()) .unwrap_or_else(|_| { dirs::home_dir() .unwrap_or_else(|| PathBuf::from(".")) diff --git a/src/crates/assembly/core/src/miniapp/builtin/mod.rs b/src/crates/assembly/core/src/miniapp/builtin/mod.rs index 113e4475a..e961bc48c 100644 --- a/src/crates/assembly/core/src/miniapp/builtin/mod.rs +++ b/src/crates/assembly/core/src/miniapp/builtin/mod.rs @@ -154,14 +154,36 @@ mod tests { MiniAppCustomizationMetadata, MiniAppCustomizationOrigin, MiniAppCustomizationOriginKind, }; - fn test_manager() -> Arc { + struct TestMiniAppManager { + manager: Arc, + root: std::path::PathBuf, + } + + impl std::ops::Deref for TestMiniAppManager { + type Target = Arc; + + fn deref(&self) -> &Self::Target { + &self.manager + } + } + + impl Drop for TestMiniAppManager { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.root); + } + } + + fn test_manager() -> TestMiniAppManager { let root = std::env::temp_dir().join(format!( "bitfun-miniapp-builtin-customization-{}", uuid::Uuid::new_v4() )); let path_manager = - Arc::new(crate::infrastructure::PathManager::with_user_root_for_tests(root)); - Arc::new(MiniAppManager::new(path_manager)) + Arc::new(crate::infrastructure::PathManager::with_user_root_for_tests(root.clone())); + TestMiniAppManager { + manager: Arc::new(MiniAppManager::new(path_manager)), + root, + } } async fn write_outdated_builtin_marker(app_dir: &std::path::Path) { diff --git a/src/crates/assembly/core/src/miniapp/manager.rs b/src/crates/assembly/core/src/miniapp/manager.rs index 00f364935..1fdde484c 100644 --- a/src/crates/assembly/core/src/miniapp/manager.rs +++ b/src/crates/assembly/core/src/miniapp/manager.rs @@ -710,14 +710,36 @@ mod tests { STORAGE_JSON, STYLE_CSS, UI_JS, WORKER_JS, }; - fn test_manager() -> MiniAppManager { + struct TestMiniAppManager { + manager: MiniAppManager, + root: std::path::PathBuf, + } + + impl std::ops::Deref for TestMiniAppManager { + type Target = MiniAppManager; + + fn deref(&self) -> &Self::Target { + &self.manager + } + } + + impl Drop for TestMiniAppManager { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.root); + } + } + + fn test_manager() -> TestMiniAppManager { let root = std::env::temp_dir().join(format!( "bitfun-miniapp-manager-draft-{}", uuid::Uuid::new_v4() )); let path_manager = - Arc::new(crate::infrastructure::PathManager::with_user_root_for_tests(root)); - MiniAppManager::new(path_manager) + Arc::new(crate::infrastructure::PathManager::with_user_root_for_tests(root.clone())); + TestMiniAppManager { + manager: MiniAppManager::new(path_manager), + root, + } } fn sample_source(css: &str) -> MiniAppSource { diff --git a/src/crates/assembly/core/src/product_runtime/runtime_services.rs b/src/crates/assembly/core/src/product_runtime/runtime_services.rs index 513521248..71d4b5fe3 100644 --- a/src/crates/assembly/core/src/product_runtime/runtime_services.rs +++ b/src/crates/assembly/core/src/product_runtime/runtime_services.rs @@ -28,7 +28,7 @@ impl CoreRuntimeServicesProvider { impl RuntimeServicesProvider for CoreRuntimeServicesProvider { fn register(&self, builder: RuntimeServicesBuilder) -> RuntimeServicesBuilder { - let session_store: Arc = Arc::new(CoreSessionStorePort); + let session_store: Arc = Arc::new(CoreSessionStorePort::default()); let builder = builder .with_session_store(session_store) .with_optional_terminal(Some(RuntimeServiceMarkerPort::terminal_port())) diff --git a/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs b/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs index 897895fc1..8e1b143a2 100644 --- a/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs +++ b/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs @@ -1293,6 +1293,7 @@ async fn create_session(state: &mut BotChatState, agent_type: &str) -> HandleRes session_name, agent_type, Some(workspace_path.clone()), + Default::default(), RemoteConnectSubmissionSource::Bot, ); let runtime = match CoreServiceAgentRuntime::agent_runtime(coordinator.clone()) { diff --git a/src/crates/assembly/core/src/service/remote_connect/remote_server.rs b/src/crates/assembly/core/src/service/remote_connect/remote_server.rs index 2bbae3bfd..f11d56acf 100644 --- a/src/crates/assembly/core/src/service/remote_connect/remote_server.rs +++ b/src/crates/assembly/core/src/service/remote_connect/remote_server.rs @@ -298,6 +298,7 @@ mod tests { use bitfun_services_integrations::remote_connect::{ remote_session_restore_target, resolve_remote_cancel_decision, resolve_remote_execution_image_contexts, RemoteCancelDecision, + RemoteDialogWorkspaceBinding, }; #[test] @@ -464,14 +465,13 @@ mod tests { #[test] fn remote_restore_target_only_restores_cold_sessions_with_workspace_binding() { + let binding = RemoteDialogWorkspaceBinding::local("/workspace/project"); + assert_eq!( - remote_session_restore_target(false, Some("/workspace/project")), + remote_session_restore_target(false, Some(&binding)), Some("/workspace/project") ); - assert_eq!( - remote_session_restore_target(true, Some("/workspace/project")), - None - ); + assert_eq!(remote_session_restore_target(true, Some(&binding)), None); assert_eq!(remote_session_restore_target(false, None), None); } @@ -505,6 +505,8 @@ mod tests { let list = serde_json::to_value(RemoteCommand::ListSessions { workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: None, + remote_ssh_host: None, limit: Some(30), offset: Some(0), query: Some("alpha".to_string()), diff --git a/src/crates/assembly/core/src/service/remote_ssh/workspace_state.rs b/src/crates/assembly/core/src/service/remote_ssh/workspace_state.rs index f1eb4efab..8bfe3f17c 100644 --- a/src/crates/assembly/core/src/service/remote_ssh/workspace_state.rs +++ b/src/crates/assembly/core/src/service/remote_ssh/workspace_state.rs @@ -5,8 +5,9 @@ //! **`(connection_id, remote_root_path)`** — *not* by remote path alone, so two //! different servers opened at the same path (e.g. `/`) do not overwrite each other. -use crate::infrastructure::{get_path_manager_arc, PathManager}; +use crate::infrastructure::get_path_manager_arc; use crate::service::remote_ssh::{RemoteFileService, RemoteTerminalManager, SSHConnectionManager}; +use crate::service::workspace_runtime::WorkspaceRuntimeService; pub use bitfun_services_integrations::remote_ssh::{ local_workspace_stable_storage_id, normalize_remote_workspace_path, remote_root_to_mirror_subpath, remote_workspace_stable_id, @@ -39,14 +40,6 @@ impl WorkspaceSessionIdentity { pub fn logical_workspace_path(&self) -> &str { &self.logical_workspace_path } - - pub fn session_storage_path(&self) -> PathBuf { - if self.is_remote() { - remote_workspace_session_mirror_dir(&self.hostname, &self.logical_workspace_path) - } else { - PathBuf::from(&self.logical_workspace_path) - } - } } /// Build a unified session identity for local or remote workspaces. @@ -122,7 +115,7 @@ pub async fn resolve_workspace_session_identity( /// Local directory where persisted sessions for this remote workspace root are stored. pub fn remote_workspace_runtime_root(ssh_host: &str, remote_root_norm: &str) -> PathBuf { bitfun_services_integrations::remote_ssh::remote_workspace_runtime_root( - PathManager::remote_ssh_mirror_root(), + get_path_manager_arc().remote_ssh_mirror_root_dir(), ssh_host, remote_root_norm, ) @@ -131,7 +124,7 @@ pub fn remote_workspace_runtime_root(ssh_host: &str, remote_root_norm: &str) -> /// Local directory where persisted sessions for this remote workspace root are stored. pub fn remote_workspace_session_mirror_dir(ssh_host: &str, remote_root_norm: &str) -> PathBuf { bitfun_services_integrations::remote_ssh::remote_workspace_session_mirror_dir( - PathManager::remote_ssh_mirror_root(), + get_path_manager_arc().remote_ssh_mirror_root_dir(), ssh_host, remote_root_norm, ) @@ -161,7 +154,7 @@ pub fn unresolved_remote_session_storage_dir( workspace_path_norm: &str, ) -> PathBuf { bitfun_services_integrations::remote_ssh::unresolved_remote_session_storage_dir( - PathManager::remote_ssh_mirror_root(), + get_path_manager_arc().remote_ssh_mirror_root_dir(), connection_id, workspace_path_norm, ) @@ -337,28 +330,35 @@ impl RemoteWorkspaceStateManager { remote_workspace_session_mirror_dir(ssh_host, remote_root_norm) } - /// Map a workspace path to the effective session storage path. - /// When `remote_connection_id` is set, remote roots map to the local session mirror dir; - /// otherwise the path is returned as-is (no path-only inference). + /// Map a workspace path to the final on-disk sessions directory. + /// Local roots map to `~/.bitfun/projects//sessions`; + /// remote roots map to the local SSH mirror sessions dir. pub async fn get_effective_session_path( &self, workspace_path: &str, remote_connection_id: Option<&str>, remote_ssh_host: Option<&str>, ) -> PathBuf { + let runtime_service = WorkspaceRuntimeService::new(get_path_manager_arc()); let remote_id = remote_connection_id .map(str::trim) .filter(|s| !s.is_empty()); if remote_id.is_none() { - return PathBuf::from(workspace_path); + return runtime_service + .context_for_local_workspace(Path::new(workspace_path)) + .sessions_dir; } let path_norm = normalize_remote_workspace_path(workspace_path); if let Some(host) = remote_ssh_host.map(str::trim).filter(|s| !s.is_empty()) { - return remote_workspace_session_mirror_dir(host, &path_norm); + return runtime_service + .context_for_remote_workspace(host, &path_norm) + .sessions_dir; } if let Some(entry) = self.lookup_connection(workspace_path, remote_id).await { if !entry.ssh_host.trim().is_empty() { - return remote_workspace_session_mirror_dir(&entry.ssh_host, &entry.remote_root); + return runtime_service + .context_for_remote_workspace(&entry.ssh_host, &entry.remote_root) + .sessions_dir; } return unresolved_remote_session_storage_dir(remote_id.unwrap(), &path_norm); } @@ -394,6 +394,7 @@ pub async fn get_effective_session_path( remote_connection_id: Option<&str>, remote_ssh_host: Option<&str>, ) -> std::path::PathBuf { + let runtime_service = WorkspaceRuntimeService::new(get_path_manager_arc()); if let Some(identity) = resolve_workspace_session_identity(workspace_path, remote_connection_id, remote_ssh_host) .await @@ -406,10 +407,19 @@ pub async fn get_effective_session_path( ); } } - return identity.session_storage_path(); + if identity.hostname == LOCAL_WORKSPACE_SSH_HOST { + return runtime_service + .context_for_local_workspace(Path::new(identity.logical_workspace_path())) + .sessions_dir; + } + return runtime_service + .context_for_remote_workspace(&identity.hostname, identity.logical_workspace_path()) + .sessions_dir; } - std::path::PathBuf::from(workspace_path) + runtime_service + .context_for_local_workspace(Path::new(workspace_path)) + .sessions_dir } /// Check if a specific path belongs to any registered remote workspace. @@ -455,7 +465,6 @@ mod tests { LOCAL_WORKSPACE_SSH_HOST, }; use crate::infrastructure::PathManager; - use std::path::PathBuf; #[tokio::test] async fn local_assistant_path_not_remote_without_connection_id() { @@ -604,7 +613,7 @@ mod tests { } #[test] - fn remote_workspace_session_identity_uses_mirror_dir_for_storage() { + fn remote_workspace_session_identity_tracks_logical_root() { let identity = workspace_session_identity( "/home/wsp/projects/test", Some("conn-1"), @@ -614,14 +623,11 @@ mod tests { assert_eq!(identity.hostname, "127.0.0.1"); assert_eq!(identity.logical_workspace_path(), "/home/wsp/projects/test"); - assert_eq!( - identity.session_storage_path(), - remote_workspace_session_mirror_dir("127.0.0.1", "/home/wsp/projects/test") - ); + assert!(identity.is_remote()); } #[test] - fn local_workspace_session_identity_uses_workspace_root_for_storage() { + fn local_workspace_session_identity_tracks_logical_root() { let workspace_root = std::env::temp_dir().join(format!( "bitfun-workspace-identity-{}", uuid::Uuid::new_v4() @@ -632,10 +638,60 @@ mod tests { .expect("local identity should resolve"); assert_eq!(identity.hostname, LOCAL_WORKSPACE_SSH_HOST); + assert!(!identity.is_remote()); + + let _ = std::fs::remove_dir_all(workspace_root); + } + + #[tokio::test] + async fn effective_session_path_returns_local_sessions_dir() { + let workspace_root = std::env::temp_dir().join(format!( + "bitfun-local-session-path-test-{}", + uuid::Uuid::new_v4() + )); + std::fs::create_dir_all(&workspace_root).expect("workspace root should exist"); + + let expected = + crate::infrastructure::get_path_manager_arc().project_sessions_dir(&workspace_root); + let actual = + super::get_effective_session_path(&workspace_root.to_string_lossy(), None, None).await; + + assert_eq!(actual, expected); + + let _ = std::fs::remove_dir_all(workspace_root); + } + + #[tokio::test] + async fn effective_session_path_returns_remote_sessions_dir() { + let actual = super::get_effective_session_path( + "/home/wsp/projects/test", + Some("conn-1"), + Some("example-host"), + ) + .await; + assert_eq!( - identity.session_storage_path(), - PathBuf::from(identity.logical_workspace_path()) + actual, + remote_workspace_session_mirror_dir("example-host", "/home/wsp/projects/test") ); + } + + #[tokio::test] + async fn manager_effective_session_path_returns_local_sessions_dir() { + let workspace_root = std::env::temp_dir().join(format!( + "bitfun-manager-local-session-path-test-{}", + uuid::Uuid::new_v4() + )); + std::fs::create_dir_all(&workspace_root).expect("workspace root should exist"); + + let manager = super::RemoteWorkspaceStateManager::new(); + let expected = + crate::infrastructure::get_path_manager_arc().project_sessions_dir(&workspace_root); + let actual = manager + .get_effective_session_path(&workspace_root.to_string_lossy(), None, None) + .await; + + assert_eq!(actual, expected); let _ = std::fs::remove_dir_all(workspace_root); } diff --git a/src/crates/assembly/core/src/service/snapshot/manager.rs b/src/crates/assembly/core/src/service/snapshot/manager.rs index 68d8497de..d7e54e541 100644 --- a/src/crates/assembly/core/src/service/snapshot/manager.rs +++ b/src/crates/assembly/core/src/service/snapshot/manager.rs @@ -801,6 +801,10 @@ mod tests { reset_snapshot_manager_new_count_for_test, set_snapshot_manager_new_delay_for_test, snapshot_manager_new_count_for_test, }; + use crate::infrastructure::PathManager; + use crate::service::workspace_runtime::{ + set_workspace_runtime_service_for_current_test, WorkspaceRuntimeService, + }; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; @@ -833,6 +837,11 @@ mod tests { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn concurrent_get_or_create_initializes_snapshot_manager_once_per_workspace() { let workspace = TestWorkspace::new(); + let _runtime_guard = set_workspace_runtime_service_for_current_test(Arc::new( + WorkspaceRuntimeService::new(Arc::new(PathManager::with_user_root_for_tests( + workspace.path().join("user-root"), + ))), + )); clear_snapshot_manager_for_test(workspace.path()); reset_snapshot_manager_new_count_for_test(); set_snapshot_manager_new_delay_for_test(Duration::from_millis(80)); diff --git a/src/crates/assembly/core/src/service/workspace_runtime/mod.rs b/src/crates/assembly/core/src/service/workspace_runtime/mod.rs index d7493ade7..3580d76f1 100644 --- a/src/crates/assembly/core/src/service/workspace_runtime/mod.rs +++ b/src/crates/assembly/core/src/service/workspace_runtime/mod.rs @@ -1,6 +1,8 @@ pub mod service; pub mod types; +#[cfg(test)] +pub use service::set_workspace_runtime_service_for_current_test; pub use service::{ get_workspace_runtime_service_arc, try_get_workspace_runtime_service_arc, WorkspaceRuntimeService, diff --git a/src/crates/assembly/core/src/service/workspace_runtime/service.rs b/src/crates/assembly/core/src/service/workspace_runtime/service.rs index db82dbad2..ee1c909ae 100644 --- a/src/crates/assembly/core/src/service/workspace_runtime/service.rs +++ b/src/crates/assembly/core/src/service/workspace_runtime/service.rs @@ -480,6 +480,11 @@ fn init_global_workspace_runtime_service() -> Arc { } pub fn get_workspace_runtime_service_arc() -> Arc { + #[cfg(test)] + if let Some(service) = test_workspace_runtime_service_override() { + return service; + } + GLOBAL_WORKSPACE_RUNTIME_SERVICE .get_or_init(init_global_workspace_runtime_service) .clone() @@ -489,6 +494,42 @@ pub fn try_get_workspace_runtime_service_arc() -> BitFunResult>> = + const { std::cell::RefCell::new(None) }; +} + +#[cfg(test)] +fn test_workspace_runtime_service_override() -> Option> { + TEST_WORKSPACE_RUNTIME_SERVICE.with(|slot| slot.borrow().clone()) +} + +#[cfg(test)] +pub struct WorkspaceRuntimeServiceOverrideGuard { + previous: Option>, +} + +#[cfg(test)] +impl Drop for WorkspaceRuntimeServiceOverrideGuard { + fn drop(&mut self) { + TEST_WORKSPACE_RUNTIME_SERVICE.with(|slot| { + *slot.borrow_mut() = self.previous.take(); + }); + } +} + +#[cfg(test)] +pub fn set_workspace_runtime_service_for_current_test( + service: Arc, +) -> WorkspaceRuntimeServiceOverrideGuard { + let previous = TEST_WORKSPACE_RUNTIME_SERVICE.with(|slot| { + let mut slot = slot.borrow_mut(); + slot.replace(service) + }); + WorkspaceRuntimeServiceOverrideGuard { previous } +} + #[cfg(test)] mod tests { use super::WorkspaceRuntimeService; diff --git a/src/crates/assembly/core/src/service_agent_runtime.rs b/src/crates/assembly/core/src/service_agent_runtime.rs index b776b21f0..05a3adf53 100644 --- a/src/crates/assembly/core/src/service_agent_runtime.rs +++ b/src/crates/assembly/core/src/service_agent_runtime.rs @@ -11,7 +11,8 @@ use bitfun_runtime_ports::{ AgentSessionCreateRequest, AgentSessionManagementPort, AgentSubmissionPort, AgentSubmissionSource, AgentTurnCancellationPort, AgentTurnCancellationRequest, RemoteControlStatePort, RemoteControlStateRequest, RemoteControlStateSnapshot, - RuntimeServiceCapability, RuntimeServicePort, + RemoteSessionWorkspaceIdentity, RuntimeServiceCapability, RuntimeServicePort, + SessionStoragePathRequest, SessionStorePort, }; use bitfun_services_integrations::remote_connect::{ build_remote_chat_messages, build_remote_model_catalog, @@ -24,14 +25,15 @@ use bitfun_services_integrations::remote_connect::{ RemoteChatHistoryToolCall, RemoteChatHistoryToolItem, RemoteChatHistoryTurn, RemoteConnectSubmissionSource, RemoteDefaultModelsConfig, RemoteDialogQueuePriority, RemoteDialogResolvedSubmission, RemoteDialogRuntimeHost, RemoteDialogSchedulerOutcomeFact, - RemoteDialogSubmissionPolicy, RemoteDialogSubmitOutcome, RemoteImageContext, - RemoteImageContextAdapter, RemoteInitialSyncRuntimeHost, RemoteInteractionRuntimeHost, - RemoteModelCapabilityFact, RemoteModelCatalog, RemoteModelCatalogFacts, RemoteModelFacts, - RemotePollRuntimeHost, RemoteReasoningModeFact, RemoteRecentWorkspaceFacts, - RemoteSessionMetadata, RemoteSessionRuntimeHost, RemoteSessionStateTracker, - RemoteSessionTrackerHost, RemoteTerminalPrewarmRequest, RemoteWorkspaceFacts, - RemoteWorkspaceFileRuntimeHost, RemoteWorkspaceKind as RemoteConnectWorkspaceKind, - RemoteWorkspaceRuntimeHost, RemoteWorkspaceUpdate, + RemoteDialogSubmissionPolicy, RemoteDialogSubmitOutcome, RemoteDialogWorkspaceBinding, + RemoteImageContext, RemoteImageContextAdapter, RemoteInitialSyncRuntimeHost, + RemoteInteractionRuntimeHost, RemoteModelCapabilityFact, RemoteModelCatalog, + RemoteModelCatalogFacts, RemoteModelFacts, RemotePollRuntimeHost, RemoteReasoningModeFact, + RemoteRecentWorkspaceFacts, RemoteSessionMetadata, RemoteSessionRuntimeHost, + RemoteSessionStateTracker, RemoteSessionTrackerHost, RemoteTerminalPrewarmRequest, + RemoteWorkspaceFacts, RemoteWorkspaceFileRuntimeHost, + RemoteWorkspaceKind as RemoteConnectWorkspaceKind, RemoteWorkspaceRuntimeHost, + RemoteWorkspaceUpdate, }; use log::{debug, error, info}; use std::sync::Arc; @@ -41,6 +43,7 @@ use crate::agentic::coordination::{ DialogScheduler, DialogSubmissionPolicy, DialogSubmitOutcome, DialogTriggerSource, }; use crate::agentic::image_analysis::ImageContextData; +use crate::agentic::session::session_store_port::CoreSessionStorePort; use crate::service::remote_connect::remote_server::RemoteExecutionDispatcher; use crate::service::config::types::{AIConfig, GlobalConfig, ModelCapability, ReasoningMode}; @@ -77,6 +80,18 @@ fn git_branch_for_workspace_path(path: &std::path::Path) -> Option { .filter(|s| !s.is_empty() && s != "HEAD") } +fn workspace_metadata_string( + metadata: &std::collections::HashMap, + key: &str, +) -> Option { + metadata + .get(key) + .and_then(|value| value.as_str()) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + async fn current_remote_workspace_facts() -> Option { let workspace_service = crate::service::workspace::get_global_workspace_service()?; workspace_service @@ -90,6 +105,11 @@ async fn current_remote_workspace_facts() -> Option { git_branch: git_branch_for_workspace_path(&root_path), kind: remote_workspace_kind(workspace.workspace_kind), assistant_id: workspace.assistant_id, + remote_connection_id: workspace_metadata_string( + &workspace.metadata, + "connectionId", + ), + remote_ssh_host: workspace_metadata_string(&workspace.metadata, "sshHost"), } }) } @@ -121,8 +141,21 @@ async fn open_workspace_with_snapshot( async fn load_remote_session_metadata_for_workspace( workspace_path: &std::path::Path, + workspace_identity: RemoteSessionWorkspaceIdentity, ) -> Result, String> { let workspace_path_display = workspace_path.to_string_lossy().to_string(); + let session_storage_dir = CoreSessionStorePort::default() + .resolve_session_storage_path(SessionStoragePathRequest { + workspace_path: workspace_path.to_path_buf(), + remote_connection_id: workspace_identity.remote_connection_id, + remote_ssh_host: workspace_identity.remote_ssh_host, + }) + .await + .map(|resolution| resolution.effective_storage_path) + .map_err(|error| { + debug!("Session storage path resolution failed for {workspace_path_display}: {error}"); + format!("Failed to resolve session storage for workspace: {error}") + })?; let path_manager = crate::infrastructure::PathManager::new() .map_err(|_| "Failed to initialize path manager".to_string())?; let path_manager = std::sync::Arc::new(path_manager); @@ -132,7 +165,7 @@ async fn load_remote_session_metadata_for_workspace( format!("Failed to initialize session storage: {error}") })?; let metadata = store - .list_session_metadata(workspace_path) + .list_session_metadata(&session_storage_dir) .await .map_err(|error| { debug!("Session list read failed for {workspace_path_display}: {error}"); @@ -363,10 +396,10 @@ async fn resolve_session_model_id(session_id: &str) -> Option { return normalize_remote_session_model_id(session.config.model_id.clone()); } - let workspace_path = - CoreServiceAgentRuntime::resolve_session_workspace_path(session_id).await?; + let session_storage_dir = + CoreServiceAgentRuntime::resolve_session_storage_dir(session_id).await?; coordinator - .restore_session(&workspace_path, session_id) + .restore_session(&session_storage_dir, session_id) .await .ok() .and_then(|session| normalize_remote_session_model_id(session.config.model_id.clone())) @@ -466,18 +499,35 @@ impl RemoteImageContextAdapter for ImageContextData { pub(crate) struct CoreServiceAgentRuntime; impl CoreServiceAgentRuntime { - pub(crate) async fn resolve_session_workspace_path( + pub(crate) async fn resolve_session_storage_dir( session_id: &str, ) -> Option { let coordinator = get_global_coordinator()?; - coordinator.resolve_session_workspace_path(session_id).await + coordinator + .get_session_manager() + .resolve_session_workspace_binding(session_id) + .await + .map(|binding| binding.session_storage_dir()) + } + + pub(crate) async fn resolve_session_logical_workspace_path( + session_id: &str, + ) -> Option { + let coordinator = get_global_coordinator()?; + coordinator + .get_session_manager() + .resolve_session_workspace_binding(session_id) + .await + .map(|binding| binding.logical_workspace_path().to_path_buf()) } pub(crate) async fn resolve_remote_file_workspace_root( session_id: Option<&str>, ) -> Option { if let Some(session_id) = session_id { - if let Some(workspace_path) = Self::resolve_session_workspace_path(session_id).await { + if let Some(workspace_path) = + Self::resolve_session_logical_workspace_path(session_id).await + { return Some(workspace_path); } } @@ -526,7 +576,7 @@ impl CoreServiceAgentRuntime { } pub(crate) async fn load_remote_chat_messages( - workspace_path: &std::path::Path, + session_storage_dir: &std::path::Path, session_id: &str, ) -> (Vec, bool) { let Ok(pm) = crate::infrastructure::PathManager::new() else { @@ -536,7 +586,10 @@ impl CoreServiceAgentRuntime { let Ok(store) = crate::agentic::persistence::PersistenceManager::new(pm) else { return (vec![], false); }; - let Ok(turns) = store.load_session_turns(workspace_path, session_id).await else { + let Ok(turns) = store + .load_session_turns(session_storage_dir, session_id) + .await + else { return (vec![], false); }; (remote_chat_messages_from_turns(&turns), false) @@ -626,14 +679,14 @@ impl CoreServiceAgentRuntime { .get_session(session_id) .is_none() { - let Some(workspace_path) = Self::resolve_session_workspace_path(session_id).await + let Some(session_storage_dir) = Self::resolve_session_storage_dir(session_id).await else { return Err(format!( - "Workspace path not available for session: {session_id}" + "Session storage directory not available for session: {session_id}" )); }; coordinator - .restore_session(&workspace_path, session_id) + .restore_session(&session_storage_dir, session_id) .await .map_err(|e| format!("Failed to restore session: {e}"))?; } @@ -938,11 +991,24 @@ impl RemoteDialogRuntimeHost for CoreRemoteDialogRuntimeHost<'_> { self.dispatcher.ensure_tracker(session_id); } - async fn resolve_binding_workspace(&self, session_id: &str) -> Option { + async fn resolve_binding_workspace( + &self, + session_id: &str, + ) -> Option { self.coordinator - .resolve_session_workspace_path(session_id) + .get_session_manager() + .resolve_session_workspace_binding(session_id) .await - .map(|path| path.to_string_lossy().into_owned()) + .map(|binding| RemoteDialogWorkspaceBinding { + workspace_path: binding.logical_workspace_path_string(), + remote_connection_id: binding.connection_id().map(ToOwned::to_owned), + remote_ssh_host: if binding.is_remote() { + Some(binding.session_identity.hostname.clone()) + .filter(|value| !value.trim().is_empty()) + } else { + None + }, + }) } async fn remote_session_exists(&self, session_id: &str) -> Result { @@ -958,8 +1024,11 @@ impl RemoteDialogRuntimeHost for CoreRemoteDialogRuntimeHost<'_> { session_id: &str, workspace_path: &str, ) -> Result<(), String> { + let restore_path = CoreServiceAgentRuntime::resolve_session_storage_dir(session_id) + .await + .unwrap_or_else(|| std::path::PathBuf::from(workspace_path)); self.coordinator - .restore_session(std::path::Path::new(workspace_path), session_id) + .restore_session(&restore_path, session_id) .await .map(|_| ()) .map_err(|e| e.to_string()) @@ -1018,6 +1087,17 @@ impl RemoteDialogRuntimeHost for CoreRemoteDialogRuntimeHost<'_> { .map(agent_input_attachment_from_image_context) .collect(); + let binding_workspace = submission.binding_workspace; + let workspace_path = binding_workspace + .as_ref() + .map(|binding| binding.workspace_path.clone()); + let remote_connection_id = binding_workspace + .as_ref() + .and_then(|binding| binding.remote_connection_id.clone()); + let remote_ssh_host = binding_workspace + .as_ref() + .and_then(|binding| binding.remote_ssh_host.clone()); + self.runtime .submit_dialog_turn(AgentDialogTurnRequest { session_id: submission.session_id, @@ -1025,9 +1105,9 @@ impl RemoteDialogRuntimeHost for CoreRemoteDialogRuntimeHost<'_> { original_message: None, turn_id: Some(submission.turn_id), agent_type: submission.resolved_agent_type, - workspace_path: submission.binding_workspace, - remote_connection_id: None, - remote_ssh_host: None, + workspace_path, + remote_connection_id, + remote_ssh_host, policy, reply_route: None, prepended_reminders: Vec::new(), @@ -1110,8 +1190,9 @@ impl RemoteInitialSyncRuntimeHost for CoreRemoteWorkspaceRuntimeHost { async fn list_session_metadata( &self, workspace_path: &std::path::Path, + workspace_identity: RemoteSessionWorkspaceIdentity, ) -> Result, String> { - load_remote_session_metadata_for_workspace(workspace_path).await + load_remote_session_metadata_for_workspace(workspace_path, workspace_identity).await } } @@ -1120,8 +1201,9 @@ impl RemoteSessionRuntimeHost for CoreRemoteSessionRuntimeHost { async fn list_session_metadata( &self, workspace_path: &std::path::Path, + workspace_identity: RemoteSessionWorkspaceIdentity, ) -> Result, String> { - load_remote_session_metadata_for_workspace(workspace_path).await + load_remote_session_metadata_for_workspace(workspace_path, workspace_identity).await } async fn resolve_default_assistant_workspace_path(&self) -> Result { @@ -1180,16 +1262,16 @@ impl RemoteSessionRuntimeHost for CoreRemoteSessionRuntimeHost { return Ok(()); } - let Some(workspace_path) = - CoreServiceAgentRuntime::resolve_session_workspace_path(session_id).await + let Some(session_storage_dir) = + CoreServiceAgentRuntime::resolve_session_storage_dir(session_id).await else { return Err(format!( - "Workspace path not available for session: {}", + "Session storage directory not available for session: {}", session_id )); }; self.coordinator - .restore_session(&workspace_path, session_id) + .restore_session(&session_storage_dir, session_id) .await .map(|_| ()) .map_err(|error| format!("Failed to restore session: {error}")) @@ -1202,25 +1284,25 @@ impl RemoteSessionRuntimeHost for CoreRemoteSessionRuntimeHost { .map_err(|error| error.to_string()) } - async fn resolve_session_workspace_path(&self, session_id: &str) -> Option { - CoreServiceAgentRuntime::resolve_session_workspace_path(session_id).await + async fn resolve_session_storage_dir(&self, session_id: &str) -> Option { + CoreServiceAgentRuntime::resolve_session_storage_dir(session_id).await } async fn load_remote_chat_messages( &self, - workspace_path: &std::path::Path, + session_storage_dir: &std::path::Path, session_id: &str, ) -> (Vec, bool) { - CoreServiceAgentRuntime::load_remote_chat_messages(workspace_path, session_id).await + CoreServiceAgentRuntime::load_remote_chat_messages(session_storage_dir, session_id).await } async fn delete_session( &self, - workspace_path: &std::path::Path, + session_storage_dir: &std::path::Path, session_id: &str, ) -> Result<(), String> { self.coordinator - .delete_session(workspace_path, session_id) + .delete_session(session_storage_dir, session_id) .await .map(|_| ()) .map_err(|error| error.to_string()) @@ -1244,16 +1326,16 @@ impl RemotePollRuntimeHost for CoreRemotePollRuntimeHost<'_> { .ok() } - async fn resolve_session_workspace_path(&self, session_id: &str) -> Option { - CoreServiceAgentRuntime::resolve_session_workspace_path(session_id).await + async fn resolve_session_storage_dir(&self, session_id: &str) -> Option { + CoreServiceAgentRuntime::resolve_session_storage_dir(session_id).await } async fn load_remote_chat_messages( &self, - workspace_path: &std::path::Path, + session_storage_dir: &std::path::Path, session_id: &str, ) -> (Vec, bool) { - CoreServiceAgentRuntime::load_remote_chat_messages(workspace_path, session_id).await + CoreServiceAgentRuntime::load_remote_chat_messages(session_storage_dir, session_id).await } } @@ -1295,9 +1377,8 @@ impl RemoteInteractionRuntimeHost for CoreRemoteInteractionRuntimeHost { #[async_trait::async_trait] impl RemoteCancelRuntimeHost for CoreRemoteCancelRuntimeHost { - async fn resolve_restore_workspace(&self, session_id: &str) -> Option { - self.coordinator - .resolve_session_workspace_path(session_id) + async fn resolve_session_storage_dir(&self, session_id: &str) -> Option { + CoreServiceAgentRuntime::resolve_session_storage_dir(session_id) .await .map(|path| path.to_string_lossy().into_owned()) } @@ -1319,10 +1400,13 @@ impl RemoteCancelRuntimeHost for CoreRemoteCancelRuntimeHost { async fn restore_remote_session( &self, session_id: &str, - workspace_path: &str, + restore_path_hint: &str, ) -> Result<(), String> { + let restore_path = CoreServiceAgentRuntime::resolve_session_storage_dir(session_id) + .await + .unwrap_or_else(|| std::path::PathBuf::from(restore_path_hint)); self.coordinator - .restore_session(std::path::Path::new(workspace_path), session_id) + .restore_session(&restore_path, session_id) .await .map(|_| ()) .map_err(|error| error.to_string()) diff --git a/src/crates/contracts/runtime-ports/src/lib.rs b/src/crates/contracts/runtime-ports/src/lib.rs index 56c2638f6..d7fa09836 100644 --- a/src/crates/contracts/runtime-ports/src/lib.rs +++ b/src/crates/contracts/runtime-ports/src/lib.rs @@ -457,6 +457,36 @@ pub struct RemoteWorkspaceFacts { pub kind: RemoteWorkspaceKind, #[serde(skip_serializing_if = "Option::is_none")] pub assistant_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_connection_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_ssh_host: Option, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct RemoteSessionWorkspaceIdentity { + pub remote_connection_id: Option, + pub remote_ssh_host: Option, +} + +impl RemoteSessionWorkspaceIdentity { + pub fn new(remote_connection_id: Option, remote_ssh_host: Option) -> Self { + Self { + remote_connection_id, + remote_ssh_host, + } + } + + pub fn from_workspace(workspace: &RemoteWorkspaceFacts) -> Self { + Self::new( + workspace.remote_connection_id.clone(), + workspace.remote_ssh_host.clone(), + ) + } + + pub fn is_empty(&self) -> bool { + self.remote_connection_id.is_none() && self.remote_ssh_host.is_none() + } } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] @@ -549,6 +579,7 @@ pub trait RemoteInitialSyncRuntimeHost: Send + Sync { async fn list_session_metadata( &self, workspace_path: &Path, + workspace_identity: RemoteSessionWorkspaceIdentity, ) -> Result, String>; } @@ -697,6 +728,10 @@ pub struct AgentBackgroundResultRequest { pub agent_type: String, #[serde(skip_serializing_if = "Option::is_none")] pub workspace_path: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_connection_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_ssh_host: Option, pub content: String, #[serde(skip_serializing_if = "Option::is_none")] pub display_content: Option, @@ -718,6 +753,10 @@ pub struct AgentThreadGoalDeliveryRequest { pub agent_type: String, #[serde(skip_serializing_if = "Option::is_none")] pub workspace_path: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_connection_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_ssh_host: Option, pub kind: AgentThreadGoalDeliveryKind, pub goal: ThreadGoal, } @@ -1225,11 +1264,6 @@ pub trait AgentSessionManagementPort: Send + Sync { async fn delete_session(&self, request: AgentSessionDeleteRequest) -> PortResult<()>; - async fn resolve_session_workspace_path( - &self, - request: AgentSessionWorkspaceRequest, - ) -> PortResult>; - async fn resolve_session_workspace_binding( &self, request: AgentSessionWorkspaceRequest, @@ -1672,6 +1706,8 @@ mod tests { git_branch: Some("main".to_string()), kind: RemoteWorkspaceKind::Remote, assistant_id: Some("assistant_1".to_string()), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), }; let session = RemoteSessionMetadata { session_id: "session_1".to_string(), @@ -1684,6 +1720,8 @@ mod tests { assert_eq!(workspace.kind.as_wire_str(), "remote"); assert_eq!(workspace.assistant_id.as_deref(), Some("assistant_1")); + assert_eq!(workspace.remote_connection_id.as_deref(), Some("conn-1")); + assert_eq!(workspace.remote_ssh_host.as_deref(), Some("host-1")); assert_eq!(session.turn_count, 3); } @@ -2000,6 +2038,8 @@ mod tests { session_id: "session_1".to_string(), agent_type: "agentic".to_string(), workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), content: "full result".to_string(), display_content: Some("short result".to_string()), metadata, @@ -2010,6 +2050,8 @@ mod tests { assert_eq!(json["sessionId"], "session_1"); assert_eq!(json["agentType"], "agentic"); assert_eq!(json["workspacePath"], "/workspace/project"); + assert_eq!(json["remoteConnectionId"], "conn-1"); + assert_eq!(json["remoteSshHost"], "host-1"); assert_eq!(json["content"], "full result"); assert_eq!(json["displayContent"], "short result"); assert_eq!(json["metadata"]["kind"], "background_result"); @@ -2021,6 +2063,8 @@ mod tests { session_id: "session_1".to_string(), agent_type: "agentic".to_string(), workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), kind: AgentThreadGoalDeliveryKind::ObjectiveUpdated, goal: ThreadGoal { goal_id: "goal_1".to_string(), @@ -2041,6 +2085,8 @@ mod tests { assert_eq!(json["sessionId"], "session_1"); assert_eq!(json["agentType"], "agentic"); assert_eq!(json["workspacePath"], "/workspace/project"); + assert_eq!(json["remoteConnectionId"], "conn-1"); + assert_eq!(json["remoteSshHost"], "host-1"); assert_eq!(json["kind"], "objective_updated"); assert_eq!(json["goal"]["goalId"], "goal_1"); } diff --git a/src/crates/execution/agent-runtime/src/runtime.rs b/src/crates/execution/agent-runtime/src/runtime.rs index 3a12f4b91..b6c0913ce 100644 --- a/src/crates/execution/agent-runtime/src/runtime.rs +++ b/src/crates/execution/agent-runtime/src/runtime.rs @@ -445,20 +445,6 @@ impl AgentRuntime { .map_err(RuntimeError::from) } - pub async fn resolve_session_workspace_path( - &self, - request: AgentSessionWorkspaceRequest, - ) -> Result, RuntimeError> { - let session_management = self - .session_management - .as_ref() - .ok_or(RuntimeError::MissingSessionManagementPort)?; - session_management - .resolve_session_workspace_path(request) - .await - .map_err(RuntimeError::from) - } - pub async fn resolve_session_workspace_binding( &self, request: AgentSessionWorkspaceRequest, @@ -637,7 +623,6 @@ mod tests { cancelled_turns: Mutex>, listed_sessions: Mutex>, deleted_sessions: Mutex>, - workspace_requests: Mutex>, workspace_binding_requests: Mutex>, resolved_agent_type: Option, } @@ -663,14 +648,6 @@ mod tests { Ok(()) } - async fn resolve_session_workspace_path( - &self, - request: AgentSessionWorkspaceRequest, - ) -> PortResult> { - self.workspace_requests.lock().unwrap().push(request); - Ok(Some("/workspace/project".to_string())) - } - async fn resolve_session_workspace_binding( &self, request: AgentSessionWorkspaceRequest, @@ -957,12 +934,6 @@ mod tests { }) .await .expect("delete session"); - let workspace_path = runtime - .resolve_session_workspace_path(AgentSessionWorkspaceRequest { - session_id: "session_1".to_string(), - }) - .await - .expect("resolve workspace"); let workspace_binding = runtime .resolve_session_workspace_binding(AgentSessionWorkspaceRequest { session_id: "session_1".to_string(), @@ -972,7 +943,6 @@ mod tests { .expect("workspace binding"); assert_eq!(sessions[0].session_id, "session_1"); - assert_eq!(workspace_path.as_deref(), Some("/workspace/project")); assert_eq!(workspace_binding.workspace_path, "/workspace/project"); assert_eq!( workspace_binding.remote_connection_id.as_deref(), @@ -980,7 +950,6 @@ mod tests { ); assert_eq!(ports.listed_sessions.lock().unwrap().len(), 1); assert_eq!(ports.deleted_sessions.lock().unwrap().len(), 1); - assert_eq!(ports.workspace_requests.lock().unwrap().len(), 1); assert_eq!(ports.workspace_binding_requests.lock().unwrap().len(), 1); } @@ -1107,6 +1076,8 @@ mod tests { session_id: "session_1".to_string(), agent_type: "agentic".to_string(), workspace_path: None, + remote_connection_id: None, + remote_ssh_host: None, content: "result".to_string(), display_content: None, metadata: serde_json::Map::new(), @@ -1157,6 +1128,8 @@ mod tests { session_id: "session_1".to_string(), agent_type: "agentic".to_string(), workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), content: "result".to_string(), display_content: Some("display".to_string()), metadata: serde_json::Map::new(), @@ -1169,6 +1142,8 @@ mod tests { session_id: "session_1".to_string(), agent_type: "agentic".to_string(), workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), kind: AgentThreadGoalDeliveryKind::Resumed, goal: ThreadGoal { goal_id: "goal_1".to_string(), @@ -1193,11 +1168,23 @@ mod tests { .as_deref(), Some("display") ); + assert_eq!( + lifecycle.background_results.lock().unwrap()[0] + .remote_connection_id + .as_deref(), + Some("conn-1") + ); assert_eq!(lifecycle.thread_goals.lock().unwrap().len(), 1); assert_eq!( lifecycle.thread_goals.lock().unwrap()[0].kind, AgentThreadGoalDeliveryKind::Resumed ); + assert_eq!( + lifecycle.thread_goals.lock().unwrap()[0] + .remote_ssh_host + .as_deref(), + Some("host-1") + ); } #[tokio::test] diff --git a/src/crates/execution/agent-runtime/src/scheduler.rs b/src/crates/execution/agent-runtime/src/scheduler.rs index 4196c9f41..74425519e 100644 --- a/src/crates/execution/agent-runtime/src/scheduler.rs +++ b/src/crates/execution/agent-runtime/src/scheduler.rs @@ -20,6 +20,8 @@ pub const DEFAULT_MAX_DIALOG_QUEUE_DEPTH: usize = 20; pub struct ActiveDialogTurn { turn_id: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, agent_type: String, user_input: String, user_message_metadata: Option, @@ -31,6 +33,8 @@ impl ActiveDialogTurn { pub fn new( turn_id: String, workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, agent_type: String, user_input: String, user_message_metadata: Option, @@ -40,6 +44,8 @@ impl ActiveDialogTurn { Self { turn_id, workspace_path, + remote_connection_id, + remote_ssh_host, agent_type, user_input, user_message_metadata, @@ -60,6 +66,22 @@ impl ActiveDialogTurn { self.workspace_path.clone() } + pub fn remote_connection_id(&self) -> Option<&str> { + self.remote_connection_id.as_deref() + } + + pub fn remote_connection_id_owned(&self) -> Option { + self.remote_connection_id.clone() + } + + pub fn remote_ssh_host(&self) -> Option<&str> { + self.remote_ssh_host.as_deref() + } + + pub fn remote_ssh_host_owned(&self) -> Option { + self.remote_ssh_host.clone() + } + pub fn agent_type(&self) -> &str { &self.agent_type } diff --git a/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs b/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs index 9afc921d0..c3f888d2b 100644 --- a/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs +++ b/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs @@ -306,6 +306,8 @@ fn active_dialog_turn_owns_agent_session_reply_suppression_facts() { let turn = ActiveDialogTurn::new( "turn-1".to_string(), Some("workspace".to_string()), + Some("target-conn".to_string()), + Some("target-host".to_string()), "agentic".to_string(), "run task".to_string(), Some(serde_json::json!({"kind": "session_message"})), @@ -316,6 +318,8 @@ fn active_dialog_turn_owns_agent_session_reply_suppression_facts() { assert!(turn.is_agent_session_request()); assert_eq!(turn.turn_id(), "turn-1"); assert_eq!(turn.workspace_path(), Some("workspace")); + assert_eq!(turn.remote_connection_id(), Some("target-conn")); + assert_eq!(turn.remote_ssh_host(), Some("target-host")); assert_eq!(turn.agent_type(), "agentic"); assert_eq!(turn.user_input(), "run task"); assert!(turn.user_message_metadata().is_some()); @@ -329,6 +333,8 @@ fn active_dialog_turn_does_not_suppress_non_agent_session_turns() { let turn = ActiveDialogTurn::new( "turn-1".to_string(), None, + None, + None, "agentic".to_string(), "user task".to_string(), None, @@ -434,6 +440,8 @@ fn agent_session_reply_action_ignores_non_agent_session_turns() { let turn = ActiveDialogTurn::new( "turn-1".to_string(), Some("workspace".to_string()), + None, + None, "agentic".to_string(), "user task".to_string(), None, @@ -606,6 +614,8 @@ fn agent_session_turn(source_session_id: &str) -> ActiveDialogTurn { ActiveDialogTurn::new( "turn-1".to_string(), Some("workspace".to_string()), + Some("target-conn".to_string()), + Some("target-host".to_string()), "agentic".to_string(), "run task".to_string(), Some(serde_json::json!({"kind": "session_message"})), diff --git a/src/crates/execution/runtime-services/src/test_support.rs b/src/crates/execution/runtime-services/src/test_support.rs index 8a2f2c52e..fec3e08a6 100644 --- a/src/crates/execution/runtime-services/src/test_support.rs +++ b/src/crates/execution/runtime-services/src/test_support.rs @@ -64,6 +64,8 @@ impl RemoteWorkspaceRuntimeHost for FakeRuntimePort { git_branch: Some("main".to_string()), kind: RemoteWorkspaceKind::Remote, assistant_id: None, + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), }) } diff --git a/src/crates/services/services-integrations/src/remote_connect.rs b/src/crates/services/services-integrations/src/remote_connect.rs index a5cad2e7f..e6fd0b3d6 100644 --- a/src/crates/services/services-integrations/src/remote_connect.rs +++ b/src/crates/services/services-integrations/src/remote_connect.rs @@ -21,10 +21,10 @@ use bitfun_runtime_ports::{ }; pub use bitfun_runtime_ports::{ RemoteAssistantWorkspaceFacts, RemoteFileChunkRange, RemoteInitialSyncRuntimeHost, - RemoteProjectionPort, RemoteRecentWorkspaceFacts, RemoteSessionMetadata, RemoteWorkspaceFacts, - RemoteWorkspaceFileChunk, RemoteWorkspaceFileContent, RemoteWorkspaceFileInfo, - RemoteWorkspaceFileRuntimeHost, RemoteWorkspaceKind, RemoteWorkspacePort, - RemoteWorkspaceRuntimeHost, RemoteWorkspaceUpdate, + RemoteProjectionPort, RemoteRecentWorkspaceFacts, RemoteSessionMetadata, + RemoteSessionWorkspaceIdentity, RemoteWorkspaceFacts, RemoteWorkspaceFileChunk, + RemoteWorkspaceFileContent, RemoteWorkspaceFileInfo, RemoteWorkspaceFileRuntimeHost, + RemoteWorkspaceKind, RemoteWorkspacePort, RemoteWorkspaceRuntimeHost, RemoteWorkspaceUpdate, }; pub use device::DeviceIdentity; pub use encryption::{decrypt_from_base64, encrypt_to_base64, KeyPair}; @@ -67,6 +67,7 @@ pub fn build_remote_session_create_request( session_name: impl Into, agent_type: impl Into, workspace_path: Option>, + workspace_identity: RemoteSessionWorkspaceIdentity, source: RemoteConnectSubmissionSource, ) -> AgentSessionCreateRequest { let mut metadata = serde_json::Map::new(); @@ -79,8 +80,8 @@ pub fn build_remote_session_create_request( session_name: session_name.into(), agent_type: agent_type.into(), workspace_path: workspace_path.map(Into::into), - remote_connection_id: None, - remote_ssh_host: None, + remote_connection_id: workspace_identity.remote_connection_id, + remote_ssh_host: workspace_identity.remote_ssh_host, metadata, } } @@ -190,12 +191,12 @@ pub fn resolve_remote_execution_image_contexts( pub fn remote_session_restore_target<'a>( session_exists: bool, - binding_workspace: Option<&'a str>, + binding_workspace: Option<&'a RemoteDialogWorkspaceBinding>, ) -> Option<&'a str> { if session_exists { None } else { - binding_workspace + binding_workspace.map(|binding| binding.workspace_path.as_str()) } } @@ -231,7 +232,7 @@ pub struct RemoteCancelTaskRequest { #[async_trait::async_trait] pub trait RemoteCancelRuntimeHost: Send + Sync { - async fn resolve_restore_workspace(&self, session_id: &str) -> Option; + async fn resolve_session_storage_dir(&self, session_id: &str) -> Option; async fn remote_control_state( &self, @@ -241,7 +242,7 @@ pub trait RemoteCancelRuntimeHost: Send + Sync { async fn restore_remote_session( &self, session_id: &str, - workspace_path: &str, + restore_path_hint: &str, ) -> Result<(), String>; async fn cancel_remote_turn(&self, session_id: &str, turn_id: &str) -> Result<(), String>; @@ -258,11 +259,16 @@ where let mut state = host.remote_control_state(&session_id).await?; if state.is_none() { - let workspace_path = host - .resolve_restore_workspace(&session_id) + let session_storage_dir = host + .resolve_session_storage_dir(&session_id) .await - .ok_or_else(|| format!("Workspace path not available for session: {}", session_id))?; - host.restore_remote_session(&session_id, &workspace_path) + .ok_or_else(|| { + format!( + "Session storage directory not available for session: {}", + session_id + ) + })?; + host.restore_remote_session(&session_id, &session_storage_dir) .await .map_err(|error| format!("Session not found: {error}"))?; state = host.remote_control_state(&session_id).await?; @@ -324,12 +330,29 @@ pub struct RemoteTerminalPrewarmRequest { pub binding_workspace: Option, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RemoteDialogWorkspaceBinding { + pub workspace_path: String, + pub remote_connection_id: Option, + pub remote_ssh_host: Option, +} + +impl RemoteDialogWorkspaceBinding { + pub fn local(workspace_path: impl Into) -> Self { + Self { + workspace_path: workspace_path.into(), + remote_connection_id: None, + remote_ssh_host: None, + } + } +} + #[derive(Debug, Clone, PartialEq)] pub struct RemoteDialogResolvedSubmission { pub session_id: String, pub content: String, pub resolved_agent_type: String, - pub binding_workspace: Option, + pub binding_workspace: Option, pub image_contexts: Vec, pub policy: RemoteDialogSubmissionPolicy, pub turn_id: String, @@ -379,7 +402,10 @@ pub trait RemoteDialogRuntimeHost: Send + Sync { fn ensure_tracker(&self, session_id: &str); - async fn resolve_binding_workspace(&self, session_id: &str) -> Option; + async fn resolve_binding_workspace( + &self, + session_id: &str, + ) -> Option; async fn remote_session_exists(&self, session_id: &str) -> Result; @@ -421,7 +447,7 @@ where let session_exists = host.remote_session_exists(&session_id).await?; if let Some(workspace_path) = - remote_session_restore_target(session_exists, binding_workspace.as_deref()) + remote_session_restore_target(session_exists, binding_workspace.as_ref()) { let _ = host .restore_remote_session(&session_id, workspace_path) @@ -430,7 +456,9 @@ where host.prewarm_remote_terminal(RemoteTerminalPrewarmRequest { session_id: session_id.clone(), - binding_workspace: binding_workspace.clone(), + binding_workspace: binding_workspace + .as_ref() + .map(|binding| binding.workspace_path.clone()), }); let resolved_agent_type = resolve_remote_agent_type(agent_type.as_deref()).to_string(); @@ -804,6 +832,8 @@ pub fn remote_workspace_info_response(workspace: Option) - git_branch: workspace.git_branch, workspace_kind: Some(workspace.kind.as_wire_str().to_string()), assistant_id: workspace.assistant_id, + remote_connection_id: workspace.remote_connection_id, + remote_ssh_host: workspace.remote_ssh_host, }, None => RemoteResponse::WorkspaceInfo { has_workspace: false, @@ -812,6 +842,8 @@ pub fn remote_workspace_info_response(workspace: Option) - git_branch: None, workspace_kind: None, assistant_id: None, + remote_connection_id: None, + remote_ssh_host: None, }, } } @@ -929,18 +961,28 @@ pub fn remote_initial_sync_response( has_more_sessions: bool, authenticated_user_id: Option, ) -> RemoteResponse { - let (has_workspace, path, project_name, git_branch, workspace_kind, assistant_id) = - match workspace { - Some(workspace) => ( - true, - Some(workspace.path.clone()), - Some(workspace.name.clone()), - workspace.git_branch.clone(), - Some(workspace.kind.as_wire_str().to_string()), - workspace.assistant_id.clone(), - ), - None => (false, None, None, None, None, None), - }; + let ( + has_workspace, + path, + project_name, + git_branch, + workspace_kind, + assistant_id, + remote_connection_id, + remote_ssh_host, + ) = match workspace { + Some(workspace) => ( + true, + Some(workspace.path.clone()), + Some(workspace.name.clone()), + workspace.git_branch.clone(), + Some(workspace.kind.as_wire_str().to_string()), + workspace.assistant_id.clone(), + workspace.remote_connection_id.clone(), + workspace.remote_ssh_host.clone(), + ), + None => (false, None, None, None, None, None, None, None), + }; let workspace_path = path.as_deref(); let sessions = metadata .iter() @@ -954,6 +996,8 @@ pub fn remote_initial_sync_response( git_branch, workspace_kind, assistant_id, + remote_connection_id, + remote_ssh_host, sessions, has_more_sessions, authenticated_user_id, @@ -1002,8 +1046,13 @@ where .and_then(|path| path.file_name()) .map(|name| name.to_string_lossy().to_string()); + let workspace_identity = workspace + .as_ref() + .map(RemoteSessionWorkspaceIdentity::from_workspace) + .unwrap_or_default(); + let (sessions, has_more) = if let Some(path) = workspace_path.as_deref() { - match host.list_session_metadata(path).await { + match host.list_session_metadata(path, workspace_identity).await { Ok(metadata) => { let total = metadata.len(); let page_size = 100usize; @@ -1066,6 +1115,7 @@ pub trait RemoteSessionRuntimeHost: Send + Sync { async fn list_session_metadata( &self, workspace_path: &Path, + workspace_identity: RemoteSessionWorkspaceIdentity, ) -> Result, String>; async fn resolve_default_assistant_workspace_path(&self) -> Result; async fn create_session(&self, request: AgentSessionCreateRequest) -> Result; @@ -1080,13 +1130,17 @@ pub trait RemoteSessionRuntimeHost: Send + Sync { ) -> Result; async fn ensure_session_loaded(&self, session_id: &str) -> Result<(), String>; async fn update_session_title(&self, session_id: &str, title: &str) -> Result; - async fn resolve_session_workspace_path(&self, session_id: &str) -> Option; + async fn resolve_session_storage_dir(&self, session_id: &str) -> Option; async fn load_remote_chat_messages( &self, - workspace_path: &Path, + session_storage_dir: &Path, session_id: &str, ) -> (Vec, bool); - async fn delete_session(&self, workspace_path: &Path, session_id: &str) -> Result<(), String>; + async fn delete_session( + &self, + session_storage_dir: &Path, + session_id: &str, + ) -> Result<(), String>; fn remove_tracker(&self, session_id: &str); } @@ -1097,6 +1151,8 @@ where match command { RemoteCommand::ListSessions { workspace_path, + remote_connection_id, + remote_ssh_host, limit, offset, query, @@ -1119,7 +1175,15 @@ where .file_name() .map(|name| name.to_string_lossy().to_string()); - match host.list_session_metadata(&workspace_path).await { + let workspace_identity = RemoteSessionWorkspaceIdentity::new( + remote_connection_id.clone(), + remote_ssh_host.clone(), + ); + + match host + .list_session_metadata(&workspace_path, workspace_identity) + .await + { Ok(metadata) => { let query = query .as_deref() @@ -1149,6 +1213,8 @@ where agent_type, session_name, workspace_path, + remote_connection_id, + remote_ssh_host, } => { let agent = resolve_remote_agent_type(agent_type.as_deref()); let is_claw = agent == "Claw"; @@ -1187,6 +1253,10 @@ where session_name, agent, Some(binding_workspace), + RemoteSessionWorkspaceIdentity::new( + remote_connection_id.clone(), + remote_ssh_host.clone(), + ), RemoteConnectSubmissionSource::Relay, ); match host.create_session(request).await { @@ -1227,24 +1297,32 @@ where limit: _, before_message_id: _, } => { - let Some(workspace_path) = host.resolve_session_workspace_path(session_id).await else { + let Some(session_storage_dir) = host.resolve_session_storage_dir(session_id).await + else { return RemoteResponse::Error { - message: format!("Workspace path not available for session: {}", session_id), + message: format!( + "Session storage directory not available for session: {}", + session_id + ), }; }; let (chat_messages, has_more) = host - .load_remote_chat_messages(&workspace_path, session_id) + .load_remote_chat_messages(&session_storage_dir, session_id) .await; remote_messages_response(session_id.clone(), chat_messages, has_more) } RemoteCommand::DeleteSession { session_id } => { - let Some(workspace_path) = host.resolve_session_workspace_path(session_id).await else { + let Some(session_storage_dir) = host.resolve_session_storage_dir(session_id).await + else { return RemoteResponse::Error { - message: format!("Workspace path not available for session: {}", session_id), + message: format!( + "Session storage directory not available for session: {}", + session_id + ), }; }; - match host.delete_session(&workspace_path, session_id).await { + match host.delete_session(&session_storage_dir, session_id).await { Ok(()) => { host.remove_tracker(session_id); remote_session_deleted_response(session_id.clone()) @@ -1262,10 +1340,10 @@ where pub trait RemotePollRuntimeHost: Send + Sync { fn ensure_tracker(&self, session_id: &str) -> Arc; async fn load_model_catalog(&self, session_id: &str) -> Option; - async fn resolve_session_workspace_path(&self, session_id: &str) -> Option; + async fn resolve_session_storage_dir(&self, session_id: &str) -> Option; async fn load_remote_chat_messages( &self, - workspace_path: &Path, + session_storage_dir: &Path, session_id: &str, ) -> (Vec, bool); } @@ -1305,13 +1383,16 @@ where ); } - let Some(workspace_path) = host.resolve_session_workspace_path(session_id).await else { + let Some(session_storage_dir) = host.resolve_session_storage_dir(session_id).await else { return RemoteResponse::Error { - message: format!("Workspace path not available for session: {}", session_id), + message: format!( + "Session storage directory not available for session: {}", + session_id + ), }; }; let (all_chat_messages, _) = host - .load_remote_chat_messages(&workspace_path, session_id) + .load_remote_chat_messages(&session_storage_dir, session_id) .await; let total_msg_count = all_chat_messages.len(); let new_messages = all_chat_messages @@ -1913,6 +1994,10 @@ pub enum RemoteCommand { }, ListSessions { workspace_path: Option, + #[serde(default)] + remote_connection_id: Option, + #[serde(default)] + remote_ssh_host: Option, limit: Option, offset: Option, query: Option, @@ -1921,6 +2006,10 @@ pub enum RemoteCommand { agent_type: Option, session_name: Option, workspace_path: Option, + #[serde(default)] + remote_connection_id: Option, + #[serde(default)] + remote_ssh_host: Option, }, GetModelCatalog { session_id: Option, @@ -2004,6 +2093,10 @@ pub enum RemoteResponse { workspace_kind: Option, #[serde(skip_serializing_if = "Option::is_none")] assistant_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + remote_connection_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + remote_ssh_host: Option, }, RecentWorkspaces { workspaces: Vec, @@ -2068,6 +2161,10 @@ pub enum RemoteResponse { workspace_kind: Option, #[serde(skip_serializing_if = "Option::is_none")] assistant_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + remote_connection_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + remote_ssh_host: Option, sessions: Vec, has_more_sessions: bool, #[serde(skip_serializing_if = "Option::is_none")] @@ -3073,6 +3170,8 @@ mod tests { git_branch: Some("main".to_string()), kind: RemoteWorkspaceKind::Normal, assistant_id: None, + remote_connection_id: None, + remote_ssh_host: None, }) } @@ -3124,6 +3223,8 @@ mod tests { git_branch: Some("main".to_string()), workspace_kind: Some("normal".to_string()), assistant_id: None, + remote_connection_id: None, + remote_ssh_host: None, } ); @@ -3147,6 +3248,7 @@ mod tests { #[derive(Default)] struct FakeSessionHost { created_requests: Mutex>, + list_identities: Mutex>, removed_trackers: Mutex>, } @@ -3155,7 +3257,12 @@ mod tests { async fn list_session_metadata( &self, _workspace_path: &Path, + workspace_identity: RemoteSessionWorkspaceIdentity, ) -> Result, String> { + self.list_identities + .lock() + .unwrap() + .push(workspace_identity); Ok(vec![ RemoteSessionMetadata { session_id: "session-a".to_string(), @@ -3220,13 +3327,13 @@ mod tests { Ok(title.trim().to_string()) } - async fn resolve_session_workspace_path(&self, _session_id: &str) -> Option { + async fn resolve_session_storage_dir(&self, _session_id: &str) -> Option { Some(PathBuf::from("/workspace/project")) } async fn load_remote_chat_messages( &self, - _workspace_path: &Path, + _session_storage_dir: &Path, _session_id: &str, ) -> (Vec, bool) { ( @@ -3247,7 +3354,7 @@ mod tests { async fn delete_session( &self, - _workspace_path: &Path, + _session_storage_dir: &Path, _session_id: &str, ) -> Result<(), String> { Ok(()) @@ -3269,6 +3376,8 @@ mod tests { &host, &RemoteCommand::ListSessions { workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), limit: Some(20), offset: Some(0), query: Some("keep".to_string()), @@ -3285,6 +3394,16 @@ mod tests { sessions[0].workspace_path.as_deref(), Some("/workspace/project") ); + let list_identities = host.list_identities.lock().unwrap(); + assert_eq!( + list_identities[0].remote_connection_id.as_deref(), + Some("conn-1") + ); + assert_eq!( + list_identities[0].remote_ssh_host.as_deref(), + Some("host-1") + ); + drop(list_identities); let created = handle_remote_session_command( &host, @@ -3292,6 +3411,8 @@ mod tests { agent_type: Some("Cowork".to_string()), session_name: None, workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), }, ) .await; @@ -3308,6 +3429,14 @@ mod tests { created_requests[0].workspace_path.as_deref(), Some("/workspace/project") ); + assert_eq!( + created_requests[0].remote_connection_id.as_deref(), + Some("conn-1") + ); + assert_eq!( + created_requests[0].remote_ssh_host.as_deref(), + Some("host-1") + ); } #[tokio::test] @@ -3348,13 +3477,13 @@ mod tests { None } - async fn resolve_session_workspace_path(&self, _session_id: &str) -> Option { + async fn resolve_session_storage_dir(&self, _session_id: &str) -> Option { None } async fn load_remote_chat_messages( &self, - _workspace_path: &Path, + _session_storage_dir: &Path, _session_id: &str, ) -> (Vec, bool) { (Vec::new(), false) @@ -3381,7 +3510,8 @@ mod tests { assert_eq!( response, RemoteResponse::Error { - message: "Workspace path not available for session: session-a".to_string(), + message: "Session storage directory not available for session: session-a" + .to_string(), } ); } diff --git a/src/crates/services/services-integrations/tests/remote_connect_contracts.rs b/src/crates/services/services-integrations/tests/remote_connect_contracts.rs index 5bec1c93c..9b5104f4d 100644 --- a/src/crates/services/services-integrations/tests/remote_connect_contracts.rs +++ b/src/crates/services/services-integrations/tests/remote_connect_contracts.rs @@ -33,14 +33,15 @@ use bitfun_services_integrations::remote_connect::{ RemoteChatHistoryTurn, RemoteCommand, RemoteCommandRuntimeHost, RemoteConnectSubmissionSource, RemoteDefaultModelsConfig, RemoteDialogQueuePriority, RemoteDialogResolvedSubmission, RemoteDialogRuntimeHost, RemoteDialogSchedulerOutcomeFact, RemoteDialogSubmissionPolicy, - RemoteDialogSubmissionRequest, RemoteDialogSubmitOutcome, RemoteImageContext, - RemoteImageContextAdapter, RemoteModelCapabilityFact, RemoteModelCatalog, + RemoteDialogSubmissionRequest, RemoteDialogSubmitOutcome, RemoteDialogWorkspaceBinding, + RemoteImageContext, RemoteImageContextAdapter, RemoteModelCapabilityFact, RemoteModelCatalog, RemoteModelCatalogFacts, RemoteModelConfig, RemoteModelFacts, RemoteReasoningModeFact, RemoteRecentWorkspaceFacts, RemoteResponse, RemoteSessionMetadata, RemoteSessionStateTracker, - RemoteSessionTrackerHost, RemoteSessionTrackerRegistry, RemoteTerminalPrewarmRequest, - RemoteToolStatus, RemoteWorkspaceFacts, RemoteWorkspaceFileChunk, RemoteWorkspaceFileContent, - RemoteWorkspaceFileInfo, RemoteWorkspaceFileRuntimeHost, RemoteWorkspaceKind, - RemoteWorkspaceUpdate, TrackerEvent, REMOTE_FILE_MAX_CHUNK_BYTES, REMOTE_FILE_MAX_READ_BYTES, + RemoteSessionTrackerHost, RemoteSessionTrackerRegistry, RemoteSessionWorkspaceIdentity, + RemoteTerminalPrewarmRequest, RemoteToolStatus, RemoteWorkspaceFacts, RemoteWorkspaceFileChunk, + RemoteWorkspaceFileContent, RemoteWorkspaceFileInfo, RemoteWorkspaceFileRuntimeHost, + RemoteWorkspaceKind, RemoteWorkspaceUpdate, TrackerEvent, REMOTE_FILE_MAX_CHUNK_BYTES, + REMOTE_FILE_MAX_READ_BYTES, }; use std::path::PathBuf; use std::sync::{Arc, Mutex}; @@ -331,14 +332,12 @@ fn remote_chat_history_assembly_skips_in_progress_assistant_history() { #[test] fn remote_connect_cancel_and_restore_policy_preserve_runtime_decisions() { + let binding = RemoteDialogWorkspaceBinding::local("D:/workspace/project"); assert_eq!( - remote_session_restore_target(false, Some("D:/workspace/project")), + remote_session_restore_target(false, Some(&binding)), Some("D:/workspace/project") ); - assert_eq!( - remote_session_restore_target(true, Some("D:/workspace/project")), - None - ); + assert_eq!(remote_session_restore_target(true, Some(&binding)), None); assert_eq!(remote_session_restore_target(false, None), None); assert_eq!( @@ -415,7 +414,7 @@ fn remote_history_contract_turn(is_in_progress: bool) -> RemoteChatHistoryTurn { struct RecordingDialogHost { session_exists: bool, - binding_workspace: Option, + binding_workspace: Option, generated_turn_id: String, restore_error: bool, submit_outcome: RemoteDialogSubmitOutcome, @@ -427,7 +426,7 @@ impl RecordingDialogHost { fn new(session_exists: bool, binding_workspace: Option<&str>) -> Self { Self { session_exists, - binding_workspace: binding_workspace.map(ToOwned::to_owned), + binding_workspace: binding_workspace.map(RemoteDialogWorkspaceBinding::local), generated_turn_id: "turn-generated".to_string(), restore_error: false, submit_outcome: RemoteDialogSubmitOutcome::Started { @@ -449,6 +448,20 @@ impl RecordingDialogHost { self } + fn with_remote_binding( + mut self, + workspace_path: &str, + remote_connection_id: &str, + remote_ssh_host: &str, + ) -> Self { + self.binding_workspace = Some(RemoteDialogWorkspaceBinding { + workspace_path: workspace_path.to_string(), + remote_connection_id: Some(remote_connection_id.to_string()), + remote_ssh_host: Some(remote_ssh_host.to_string()), + }); + self + } + fn events(&self) -> Vec { self.events.lock().unwrap().clone() } @@ -473,7 +486,10 @@ impl RemoteDialogRuntimeHost for RecordingDialogHost { .push(format!("ensure_tracker:{session_id}")); } - async fn resolve_binding_workspace(&self, session_id: &str) -> Option { + async fn resolve_binding_workspace( + &self, + session_id: &str, + ) -> Option { self.events .lock() .unwrap() @@ -587,7 +603,7 @@ fn remote_state( #[async_trait::async_trait] impl RemoteCancelRuntimeHost for RecordingCancelHost { - async fn resolve_restore_workspace(&self, session_id: &str) -> Option { + async fn resolve_session_storage_dir(&self, session_id: &str) -> Option { self.events .lock() .unwrap() @@ -687,6 +703,8 @@ impl RemoteCommandRuntimeHost for RecordingCommandHost { git_branch: None, workspace_kind: None, assistant_id: None, + remote_connection_id: None, + remote_ssh_host: None, } } @@ -939,7 +957,10 @@ async fn remote_connect_dialog_runtime_owns_restore_prewarm_and_submit_order() { assert_eq!(submitted.content, "hello"); assert_eq!(submitted.resolved_agent_type, "agentic"); assert_eq!( - submitted.binding_workspace.as_deref(), + submitted + .binding_workspace + .as_ref() + .map(|binding| binding.workspace_path.as_str()), Some("D:/workspace/project") ); assert_eq!(submitted.image_contexts, vec!["image-1".to_string()]); @@ -955,6 +976,38 @@ async fn remote_connect_dialog_runtime_owns_restore_prewarm_and_submit_order() { assert!(submitted.policy.skip_tool_confirmation); } +#[tokio::test] +async fn remote_connect_dialog_runtime_preserves_remote_workspace_identity() { + let host = RecordingDialogHost::new(false, None).with_remote_binding( + "/home/wsp/project", + "ssh-1", + "dev-host", + ); + + submit_remote_dialog( + &host, + RemoteDialogSubmissionRequest { + session_id: "session-1".to_string(), + content: "hello".to_string(), + agent_type: Some("code".to_string()), + image_contexts: Vec::::new(), + policy: RemoteDialogSubmissionPolicy::for_source(RemoteConnectSubmissionSource::Relay), + turn_id: Some("turn-remote".to_string()), + }, + ) + .await + .expect("dialog submit succeeds"); + + let submitted = host.submitted(); + let binding = submitted + .binding_workspace + .as_ref() + .expect("binding workspace should be preserved"); + assert_eq!(binding.workspace_path, "/home/wsp/project"); + assert_eq!(binding.remote_connection_id.as_deref(), Some("ssh-1")); + assert_eq!(binding.remote_ssh_host.as_deref(), Some("dev-host")); +} + #[tokio::test] async fn remote_connect_dialog_runtime_preserves_explicit_turn_without_restore() { let host = RecordingDialogHost::new(true, Some("D:/workspace/project")).with_submit_outcome( @@ -1437,6 +1490,8 @@ fn remote_connect_workspace_response_helpers_own_wire_shape() { git_branch: Some("main".to_string()), kind: RemoteWorkspaceKind::Remote, assistant_id: Some("assistant-1".to_string()), + remote_connection_id: Some("ssh-1".to_string()), + remote_ssh_host: Some("dev-host".to_string()), }; let info_json = serde_json::to_value(remote_workspace_info_response(Some(workspace.clone()))) @@ -1448,6 +1503,8 @@ fn remote_connect_workspace_response_helpers_own_wire_shape() { assert_eq!(info_json["git_branch"], "main"); assert_eq!(info_json["workspace_kind"], "remote"); assert_eq!(info_json["assistant_id"], "assistant-1"); + assert_eq!(info_json["remote_connection_id"], "ssh-1"); + assert_eq!(info_json["remote_ssh_host"], "dev-host"); let empty_json = serde_json::to_value(remote_workspace_info_response(None)).expect("serialize empty info"); @@ -1569,6 +1626,8 @@ fn remote_connect_session_response_helpers_own_pagination_and_timestamps() { git_branch: Some("main".to_string()), kind: RemoteWorkspaceKind::Normal, assistant_id: None, + remote_connection_id: None, + remote_ssh_host: None, }), metadata, Some("project"), @@ -1579,6 +1638,8 @@ fn remote_connect_session_response_helpers_own_pagination_and_timestamps() { assert_eq!(initial_json["resp"], "initial_sync"); assert_eq!(initial_json["has_workspace"], true); assert_eq!(initial_json["workspace_kind"], "normal"); + assert!(initial_json.get("remote_connection_id").is_none()); + assert!(initial_json.get("remote_ssh_host").is_none()); assert_eq!(initial_json["has_more_sessions"], true); assert_eq!(initial_json["sessions"].as_array().unwrap().len(), 3); assert_eq!(initial_json["authenticated_user_id"], "user-1"); @@ -1618,6 +1679,10 @@ fn remote_connect_session_create_contract_preserves_workspace_binding() { "Remote Session", "agentic", Some("D:/workspace/project"), + RemoteSessionWorkspaceIdentity::new( + Some("ssh-1".to_string()), + Some("dev-host".to_string()), + ), RemoteConnectSubmissionSource::Relay, ); @@ -1628,6 +1693,8 @@ fn remote_connect_session_create_contract_preserves_workspace_binding() { Some("D:/workspace/project") ); assert_eq!(request.metadata["source"], "remote_relay"); + assert_eq!(request.remote_connection_id.as_deref(), Some("ssh-1")); + assert_eq!(request.remote_ssh_host.as_deref(), Some("dev-host")); } #[test] @@ -1728,12 +1795,16 @@ fn remote_connect_command_wire_shape_lives_in_owner_contract() { let list = serde_json::to_value(RemoteCommand::ListSessions { workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: Some("conn-1".to_string()), + remote_ssh_host: Some("host-1".to_string()), limit: Some(30), offset: Some(0), query: Some("alpha".to_string()), }) .expect("serialize list command"); assert_eq!(list["cmd"], "list_sessions"); + assert_eq!(list["remote_connection_id"], "conn-1"); + assert_eq!(list["remote_ssh_host"], "host-1"); assert_eq!(list["query"], "alpha"); let rename = serde_json::to_value(RemoteCommand::UpdateSessionTitle { diff --git a/src/web-ui/src/app/startup/startupPerformanceContract.test.ts b/src/web-ui/src/app/startup/startupPerformanceContract.test.ts index ee6c3d540..cdd02d067 100644 --- a/src/web-ui/src/app/startup/startupPerformanceContract.test.ts +++ b/src/web-ui/src/app/startup/startupPerformanceContract.test.ts @@ -471,7 +471,7 @@ describe('startup performance contract', () => { expect(getStart).toBeGreaterThan(-1); expect(clearStart).toBeGreaterThan(getStart); - expect(getSource).toContain('resolve_session_workspace_path_for_thread_goal_read'); + expect(getSource).toContain('resolve_thread_goal_storage_path'); expect(getSource).not.toContain('ensure_session_for_thread_goal'); expect(getSource).not.toContain('restore_session'); }); From 33708f74f3460ba429e434e244bc097b5647fe6d Mon Sep 17 00:00:00 2001 From: wsp Date: Fri, 26 Jun 2026 22:22:48 +0800 Subject: [PATCH 3/4] fix: preserve session restore path semantics Split session restore APIs by input type so logical workspace paths, resolved session storage paths, and workspace identity requests no longer share the same Path-only entrypoint. - add storage-path and workspace-request restore variants - route cron, desktop, scheduler, and remote runtime restores through the correct API - preserve remote workspace identity during remote dialog restore - keep Path-only restore methods for local/legacy callers - add regression coverage for resolved sessions dirs and remote identity restore --- src/apps/desktop/src/api/agentic_api.rs | 74 ++- .../src/agentic/coordination/coordinator.rs | 186 +++++- .../src/agentic/coordination/scheduler.rs | 2 +- .../src/agentic/session/session_manager.rs | 588 +++++++++++++++--- .../src/agentic/session/session_store_port.rs | 8 +- .../service/remote_connect/remote_server.rs | 2 +- .../core/src/service_agent_runtime.rs | 38 +- .../src/remote_connect.rs | 16 +- .../tests/remote_connect_contracts.rs | 34 +- 9 files changed, 774 insertions(+), 174 deletions(-) diff --git a/src/apps/desktop/src/api/agentic_api.rs b/src/apps/desktop/src/api/agentic_api.rs index 20d73178c..924e5511a 100644 --- a/src/apps/desktop/src/api/agentic_api.rs +++ b/src/apps/desktop/src/api/agentic_api.rs @@ -23,7 +23,7 @@ use bitfun_core::agentic::deep_review_policy::{ }; use bitfun_core::agentic::goal_mode::{ThreadGoal, ThreadGoalStatus}; use bitfun_core::agentic::image_analysis::ImageContextData; -use bitfun_core::agentic::session::{SessionViewRestoreRequest, SessionViewRestoreTiming}; +use bitfun_core::agentic::session::SessionViewRestoreTiming; use bitfun_core::agentic::tools::image_context::get_image_context; use bitfun_core::agentic::tools::implementations::exec_command::{ background_command_output_capture, control_exec_command_session, send_exec_command_input, @@ -763,7 +763,7 @@ pub async fn update_session_title( .await; coordinator - .restore_session(&effective, session_id) + .restore_session_from_storage_path(&effective, session_id) .await .map_err(|e| format!("Failed to restore session before renaming: {}", e))?; } @@ -808,10 +808,12 @@ pub async fn ensure_coordinator_session( .await; let restore_result = if request.include_internal { coordinator - .restore_internal_session(&effective, session_id) + .restore_internal_session_from_storage_path(&effective, session_id) .await } else { - coordinator.restore_session(&effective, session_id).await + coordinator + .restore_session_from_storage_path(&effective, session_id) + .await }; restore_result.map(|_| ()).map_err(|e| e.to_string()) } @@ -903,7 +905,7 @@ pub async fn compact_session( ) .await; coordinator - .restore_session(&effective, session_id) + .restore_session_from_storage_path(&effective, session_id) .await .map_err(|e| format!("Failed to restore session before compacting: {}", e))?; } @@ -951,7 +953,7 @@ pub async fn activate_session_goal( ) .await; coordinator - .restore_session(&effective, session_id) + .restore_session_from_storage_path(&effective, session_id) .await .map_err(|e| format!("Failed to restore session before activating goal mode: {e}"))?; } @@ -1001,7 +1003,7 @@ async fn ensure_session_for_thread_goal( ) .await; coordinator - .restore_session(&effective, session_id) + .restore_session_from_storage_path(&effective, session_id) .await .map_err(|e| format!("Failed to restore session before thread goal access: {e}"))?; } @@ -1220,7 +1222,7 @@ pub async fn run_init_agents_md( ) .await; coordinator - .restore_session(&effective, session_id) + .restore_session_from_storage_path(&effective, session_id) .await .map_err(|e| format!("Failed to restore session before running /init: {e}"))?; } @@ -1711,11 +1713,11 @@ pub async fn restore_session( .await; let session = if request.include_internal { coordinator - .restore_internal_session(&effective_path, &request.session_id) + .restore_internal_session_from_storage_path(&effective_path, &request.session_id) .await } else { coordinator - .restore_session(&effective_path, &request.session_id) + .restore_session_from_storage_path(&effective_path, &request.session_id) .await } .map_err(|e| format!("Failed to restore session: {}", e))?; @@ -1744,48 +1746,45 @@ pub async fn restore_session_view( request.remote_connection_id.as_deref(), request.remote_ssh_host.as_deref(), ) - .await; + .await; + let resolve_storage_path_duration_ms = + path_started_at.elapsed().as_millis().min(u64::MAX as u128) as u64; debug!( "restore_session_view storage path resolved: trace_id={}, session_id={}, duration_ms={}", trace_id, request.session_id, - path_started_at.elapsed().as_millis() + resolve_storage_path_duration_ms ); - let view_request = SessionViewRestoreRequest { - workspace_path: effective_path, - session_id: request.session_id.clone(), - include_internal: request.include_internal, - tail_turn_count: request.tail_turn_count, - }; - let tail_turn_count = view_request + let session_storage_path = effective_path; + let tail_turn_count = request .tail_turn_count .filter(|count| *count > 0) .map(|count| count.min(16)); - let (session, mut turns, total_turn_count, timings) = + let (session, mut turns, total_turn_count, mut timings) = if let Some(tail_turn_count) = tail_turn_count { - if view_request.include_internal { + if request.include_internal { coordinator - .restore_internal_session_view_tail_timed( - &view_request.workspace_path, - &view_request.session_id, + .restore_internal_session_view_from_storage_path_tail_timed( + &session_storage_path, + &request.session_id, tail_turn_count, ) .await } else { coordinator - .restore_session_view_tail_timed( - &view_request.workspace_path, - &view_request.session_id, + .restore_session_view_from_storage_path_tail_timed( + &session_storage_path, + &request.session_id, tail_turn_count, ) .await } - } else if view_request.include_internal { + } else if request.include_internal { coordinator - .restore_internal_session_view_timed( - &view_request.workspace_path, - &view_request.session_id, + .restore_internal_session_view_from_storage_path_timed( + &session_storage_path, + &request.session_id, ) .await .map(|(session, turns, timings)| { @@ -1794,7 +1793,10 @@ pub async fn restore_session_view( }) } else { coordinator - .restore_session_view_timed(&view_request.workspace_path, &view_request.session_id) + .restore_session_view_from_storage_path_timed( + &session_storage_path, + &request.session_id, + ) .await .map(|(session, turns, timings)| { let total_turn_count = turns.len(); @@ -1802,6 +1804,7 @@ pub async fn restore_session_view( }) } .map_err(|e| format!("Failed to restore session view: {}", e))?; + timings.resolve_storage_path_duration_ms = resolve_storage_path_duration_ms; let loaded_turn_count = turns.len(); let is_partial = loaded_turn_count < total_turn_count; @@ -1882,11 +1885,14 @@ pub async fn restore_session_with_turns( ); let (session, turns) = if request.include_internal { coordinator - .restore_internal_session_with_turns(&effective_path, &request.session_id) + .restore_internal_session_with_turns_from_storage_path( + &effective_path, + &request.session_id, + ) .await } else { coordinator - .restore_session_with_turns(&effective_path, &request.session_id) + .restore_session_with_turns_from_storage_path(&effective_path, &request.session_id) .await } .map_err(|e| format!("Failed to restore session: {}", e))?; diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index ce6d7a889..e483942d7 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -1707,7 +1707,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet Ok(restore_path) => { match self .session_manager - .restore_session(&restore_path, session_id) + .restore_session_from_storage_path(&restore_path, session_id) .await { Ok(_) => { @@ -2695,7 +2695,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet if needs_restore { let restore_path = self.restore_path_for_existing_session(&session_id).await?; self.session_manager - .restore_session(&restore_path, &session_id) + .restore_session_from_storage_path(&restore_path, &session_id) .await?; } @@ -2884,7 +2884,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet ) .await?; self.session_manager - .restore_session(&restore_path, &session_id) + .restore_session_from_storage_path(&restore_path, &session_id) .await? } }; @@ -3045,7 +3045,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .await?; match self .session_manager - .restore_session(&restore_path, &session_id) + .restore_session_from_storage_path(&restore_path, &session_id) .await { Ok(_) => { @@ -3875,6 +3875,46 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .await } + pub async fn restore_session_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult { + self.session_manager + .restore_session_from_storage_path(session_storage_path, session_id) + .await + } + + pub async fn restore_internal_session_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult { + self.session_manager + .restore_internal_session_from_storage_path(session_storage_path, session_id) + .await + } + + pub async fn restore_session_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult { + self.session_manager + .restore_session_for_workspace(request, session_id) + .await + } + + pub async fn restore_internal_session_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult { + self.session_manager + .restore_internal_session_for_workspace(request, session_id) + .await + } + pub async fn restore_internal_session( &self, workspace_path: &Path, @@ -3896,6 +3936,46 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .await } + pub async fn restore_session_with_turns_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + self.session_manager + .restore_session_with_turns_from_storage_path(session_storage_path, session_id) + .await + } + + pub async fn restore_internal_session_with_turns_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + self.session_manager + .restore_internal_session_with_turns_from_storage_path(session_storage_path, session_id) + .await + } + + pub async fn restore_session_with_turns_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + self.session_manager + .restore_session_with_turns_for_workspace(request, session_id) + .await + } + + pub async fn restore_internal_session_with_turns_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + self.session_manager + .restore_internal_session_with_turns_for_workspace(request, session_id) + .await + } + pub async fn restore_internal_session_with_turns( &self, workspace_path: &Path, @@ -3931,6 +4011,34 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .await } + pub async fn restore_session_view_for_workspace_timed( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult<( + Session, + Vec, + crate::agentic::session::session_manager::SessionViewRestoreTiming, + )> { + self.session_manager + .restore_session_view_for_workspace_timed(request, session_id) + .await + } + + pub async fn restore_session_view_from_storage_path_timed( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult<( + Session, + Vec, + crate::agentic::session::session_manager::SessionViewRestoreTiming, + )> { + self.session_manager + .restore_session_view_from_storage_path_timed(session_storage_path, session_id) + .await + } + pub async fn restore_session_view_tail( &self, workspace_path: &Path, @@ -3958,6 +4066,26 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .await } + pub async fn restore_session_view_from_storage_path_tail_timed( + &self, + session_storage_path: &Path, + session_id: &str, + tail_turn_count: usize, + ) -> BitFunResult<( + Session, + Vec, + usize, + crate::agentic::session::session_manager::SessionViewRestoreTiming, + )> { + self.session_manager + .restore_session_view_from_storage_path_tail_timed( + session_storage_path, + session_id, + tail_turn_count, + ) + .await + } + pub async fn restore_internal_session_view( &self, workspace_path: &Path, @@ -3982,6 +4110,34 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .await } + pub async fn restore_internal_session_view_for_workspace_timed( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult<( + Session, + Vec, + crate::agentic::session::session_manager::SessionViewRestoreTiming, + )> { + self.session_manager + .restore_internal_session_view_for_workspace_timed(request, session_id) + .await + } + + pub async fn restore_internal_session_view_from_storage_path_timed( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult<( + Session, + Vec, + crate::agentic::session::session_manager::SessionViewRestoreTiming, + )> { + self.session_manager + .restore_internal_session_view_from_storage_path_timed(session_storage_path, session_id) + .await + } + pub async fn restore_internal_session_view_tail( &self, workspace_path: &Path, @@ -4009,6 +4165,26 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .await } + pub async fn restore_internal_session_view_from_storage_path_tail_timed( + &self, + session_storage_path: &Path, + session_id: &str, + tail_turn_count: usize, + ) -> BitFunResult<( + Session, + Vec, + usize, + crate::agentic::session::session_manager::SessionViewRestoreTiming, + )> { + self.session_manager + .restore_internal_session_view_from_storage_path_tail_timed( + session_storage_path, + session_id, + tail_turn_count, + ) + .await + } + /// List all sessions pub async fn list_sessions(&self, workspace_path: &Path) -> BitFunResult> { self.session_manager.list_sessions(workspace_path).await @@ -5893,7 +6069,7 @@ impl bitfun_runtime_ports::AgentSubmissionPort for ConversationCoordinator { return Ok(None); }; - self.restore_session(&binding.session_storage_dir(), session_id) + self.restore_session_from_storage_path(&binding.session_storage_dir(), session_id) .await .map(|session| Some(session.agent_type)) .map_err(|error| { diff --git a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs index c0aac883f..29883935f 100644 --- a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs +++ b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs @@ -538,7 +538,7 @@ impl DialogScheduler { ) .await?; self.session_manager - .restore_session(&restore_path, session_id) + .restore_session_from_storage_path(&restore_path, session_id) .await .map_err(|error| error.to_string())? } diff --git a/src/crates/assembly/core/src/agentic/session/session_manager.rs b/src/crates/assembly/core/src/agentic/session/session_manager.rs index 4fdafb7bd..45a45c1a9 100644 --- a/src/crates/assembly/core/src/agentic/session/session_manager.rs +++ b/src/crates/assembly/core/src/agentic/session/session_manager.rs @@ -38,9 +38,7 @@ use crate::util::errors::{BitFunError, BitFunResult}; use crate::util::sanitize_plain_model_output; use crate::util::timing::elapsed_ms_u64; pub use bitfun_runtime_ports::SessionViewRestoreTiming; -use bitfun_runtime_ports::{ - SessionStoragePathRequest, SessionStorePort, SessionViewRestoreRequest, -}; +use bitfun_runtime_ports::{SessionStoragePathRequest, SessionStorePort}; use bitfun_services_core::session::{ apply_session_lineage, collect_hidden_subagent_cascade as collect_hidden_subagent_cascade_ids, merge_session_custom_metadata as merge_session_custom_metadata_value, @@ -115,7 +113,7 @@ pub struct SessionManager { /// or resolve workspace-bound operations that only receive a session_id. /// This cache is intentionally retained across memory eviction, but should /// be cleared when a session is explicitly deleted. - session_workspace_index: Arc>, + session_storage_path_index: Arc>, /// Sub-components context_store: Arc, @@ -464,6 +462,42 @@ impl SessionManager { .unwrap_or_else(|| workspace_path.to_path_buf()) } + async fn resolve_storage_path_for_workspace_path(&self, workspace_path: &Path) -> PathBuf { + let storage_path_started_at = Instant::now(); + let session_storage_path = self + .effective_storage_path_for_workspace_path(workspace_path) + .await; + debug!( + "Session storage path resolved from workspace: workspace_path={}, session_storage_path={}, duration_ms={}", + workspace_path.display(), + session_storage_path.display(), + elapsed_ms_u64(storage_path_started_at) + ); + session_storage_path + } + + async fn resolve_storage_path_for_request( + &self, + request: SessionStoragePathRequest, + ) -> BitFunResult { + let storage_path_started_at = Instant::now(); + let requested_workspace_path = request.workspace_path.clone(); + let session_storage_path = CoreSessionStorePort::with_path_manager( + self.persistence_manager.path_manager().clone(), + ) + .resolve_session_storage_path(request) + .await + .map(|resolution| resolution.effective_storage_path) + .map_err(|error| BitFunError::Session(error.to_string()))?; + debug!( + "Session storage path resolved from workspace request: workspace_path={}, session_storage_path={}, duration_ms={}", + requested_workspace_path.display(), + session_storage_path.display(), + elapsed_ms_u64(storage_path_started_at) + ); + Ok(session_storage_path) + } + #[allow(dead_code)] fn session_workspace_path(&self, session_id: &str) -> Option { self.sessions @@ -473,7 +507,7 @@ impl SessionManager { /// Resolve the effective storage path for a session by ID. /// For remote workspaces, maps the remote path to a local session storage path. - async fn effective_session_workspace_path(&self, session_id: &str) -> Option { + async fn effective_session_storage_path(&self, session_id: &str) -> Option { let config = self.sessions.get(session_id)?.config.clone(); self.effective_storage_path_for_config(&config).await } @@ -491,11 +525,11 @@ impl SessionManager { } } - let indexed_workspace_path = self - .session_workspace_index + let indexed_storage_path = self + .session_storage_path_index .get(session_id) .map(|entry| entry.clone()); - if let Some(session_storage_path) = indexed_workspace_path { + if let Some(session_storage_path) = indexed_storage_path { if let Some(binding) = self .resolve_persisted_session_workspace_binding( session_id, @@ -523,7 +557,7 @@ impl SessionManager { ) .await { - self.session_workspace_index + self.session_storage_path_index .insert(session_id.to_string(), session_storage_path); return Some(binding); } @@ -816,7 +850,7 @@ impl SessionManager { return; } - let Some(workspace_path) = self.effective_session_workspace_path(session_id).await else { + let Some(workspace_path) = self.effective_session_storage_path(session_id).await else { debug!( "Skipping context snapshot persistence because workspace path is unavailable: session_id={}, turn_index={}, reason={}", session_id, turn_index, reason @@ -864,7 +898,7 @@ impl SessionManager { } let cache = if self.should_persist_session_id(session_id) { - match self.effective_session_workspace_path(session_id).await { + match self.effective_session_storage_path(session_id).await { Some(workspace_path) => { match self .load_prompt_cache_from_persistence(&workspace_path, session_id) @@ -940,7 +974,7 @@ impl SessionManager { return; } - let Some(workspace_path) = self.effective_session_workspace_path(session_id).await else { + let Some(workspace_path) = self.effective_session_storage_path(session_id).await else { debug!( "Skipping prompt cache persistence because workspace path is unavailable: session_id={}, reason={}", session_id, reason @@ -986,7 +1020,7 @@ impl SessionManager { let manager = Self { sessions: Arc::new(DashMap::new()), - session_workspace_index: Arc::new(DashMap::new()), + session_storage_path_index: Arc::new(DashMap::new()), context_store, prompt_cache_store: Arc::new(SessionPromptCacheStore::new()), turn_skill_agent_snapshot_store: Arc::new(TurnSkillAgentSnapshotStore::new()), @@ -1175,7 +1209,7 @@ impl SessionManager { fn spawn_model_reconciliation_listener(&self) { let sessions = self.sessions.clone(); - let session_workspace_index = self.session_workspace_index.clone(); + let session_storage_path_index = self.session_storage_path_index.clone(); let context_store = self.context_store.clone(); let prompt_cache_store = self.prompt_cache_store.clone(); let turn_skill_agent_snapshot_store = self.turn_skill_agent_snapshot_store.clone(); @@ -1199,7 +1233,7 @@ impl SessionManager { // surface area we need from the cloned shared fields above. let manager = Self { sessions, - session_workspace_index, + session_storage_path_index, context_store, prompt_cache_store, turn_skill_agent_snapshot_store, @@ -1339,7 +1373,7 @@ impl SessionManager { // 1. Add to memory self.sessions.insert(session_id.clone(), session.clone()); - self.session_workspace_index + self.session_storage_path_index .insert(session_id.clone(), session_storage_path.clone()); // 2. Initialize the in-memory context cache. @@ -1477,7 +1511,7 @@ impl SessionManager { return None; } - let workspace_path = self.effective_session_workspace_path(session_id).await?; + let workspace_path = self.effective_session_storage_path(session_id).await?; match self .load_turn_skill_agent_snapshot_from_persistence( &workspace_path, @@ -1526,7 +1560,7 @@ impl SessionManager { return cached_snapshot; } - let workspace_path = self.effective_session_workspace_path(session_id).await?; + let workspace_path = self.effective_session_storage_path(session_id).await?; let scan_floor_exclusive = cached_snapshot.as_ref().map(|snapshot| snapshot.0); for index in (0..=turn_index).rev() { if scan_floor_exclusive.is_some_and(|floor| index <= floor) { @@ -1573,7 +1607,7 @@ impl SessionManager { return; } - let Some(workspace_path) = self.effective_session_workspace_path(session_id).await else { + let Some(workspace_path) = self.effective_session_storage_path(session_id).await else { debug!( "Skipping turn skill-agent snapshot persistence because workspace path is unavailable: session_id={}, turn_index={}", session_id, turn_index @@ -1610,7 +1644,7 @@ impl SessionManager { return; } - let Some(workspace_path) = self.effective_session_workspace_path(session_id).await else { + let Some(workspace_path) = self.effective_session_storage_path(session_id).await else { debug!( "Skipping first-turn skill-agent baseline recovery persistence because workspace path is unavailable: session_id={}", session_id @@ -1657,7 +1691,7 @@ impl SessionManager { return; } - let Some(workspace_path) = self.effective_session_workspace_path(session_id).await else { + let Some(workspace_path) = self.effective_session_storage_path(session_id).await else { debug!( "Skipping listing reminder baseline override persistence because workspace path is unavailable: session_id={}", session_id @@ -1698,7 +1732,7 @@ impl SessionManager { return None; } - let workspace_path = self.effective_session_workspace_path(session_id).await?; + let workspace_path = self.effective_session_storage_path(session_id).await?; let snapshot = match self .persistence_manager .load_skill_agent_baseline_override_snapshot(&workspace_path, session_id) @@ -2015,7 +2049,7 @@ impl SessionManager { session_id: &str, new_state: SessionState, ) -> BitFunResult<()> { - let effective_path = self.effective_session_workspace_path(session_id).await; + let effective_path = self.effective_session_storage_path(session_id).await; // IMPORTANT: keep the DashMap guard scope short -- do NOT hold it across .await. // Collect the data needed for persistence, then release the guard before doing I/O. @@ -2059,7 +2093,7 @@ impl SessionManager { expected_turn_id: &str, new_state: SessionState, ) -> BitFunResult { - let effective_path = self.effective_session_workspace_path(session_id).await; + let effective_path = self.effective_session_storage_path(session_id).await; let should_persist = if let Some(mut session) = self.sessions.get_mut(session_id) { let owns_processing_turn = matches!( @@ -2109,7 +2143,7 @@ impl SessionManager { /// Update session title (in-memory + persistence) pub async fn update_session_title(&self, session_id: &str, title: &str) -> BitFunResult<()> { let normalized_title = Self::normalize_session_title_input(title)?; - let workspace_path = self.effective_session_workspace_path(session_id).await; + let workspace_path = self.effective_session_storage_path(session_id).await; { let Some(mut session) = self.sessions.get_mut(session_id) else { @@ -2200,7 +2234,7 @@ impl SessionManager { } if self.should_persist_session_id(session_id) { - let effective_path = self.effective_session_workspace_path(session_id).await; + let effective_path = self.effective_session_storage_path(session_id).await; let session_snapshot = self.sessions.get(session_id).map(|s| s.clone()); // Ref guard released -- DashMap shard lock is free. if let (Some(workspace_path), Some(session)) = (effective_path, session_snapshot) { @@ -2240,7 +2274,7 @@ impl SessionManager { } if self.should_persist_session_id(session_id) { - let effective_path = self.effective_session_workspace_path(session_id).await; + let effective_path = self.effective_session_storage_path(session_id).await; let session_snapshot = self.sessions.get(session_id).map(|s| s.clone()); if let (Some(workspace_path), Some(session)) = (effective_path, session_snapshot) { self.persistence_manager @@ -2300,18 +2334,20 @@ impl SessionManager { let mut resolved_context_window = None; // If the session was evicted from memory (idle > 1h), try to restore it - // using the workspace path recorded when it was first created/restored. + // using the storage path recorded when it was first created/restored. if !self.sessions.contains_key(session_id) && self.config.enable_persistence { - let workspace_path = self - .session_workspace_index + let session_storage_path = self + .session_storage_path_index .get(session_id) .map(|entry| entry.clone()); - if let Some(workspace_path) = workspace_path { + if let Some(session_storage_path) = session_storage_path { debug!( "Session evicted from memory, restoring for model update: session_id={}", session_id ); - let _ = self.restore_session(&workspace_path, session_id).await; + let _ = self + .restore_session_from_storage_path(&session_storage_path, session_id) + .await; } } @@ -2331,7 +2367,7 @@ impl SessionManager { } if self.should_persist_session_id(session_id) { - let effective_path = self.effective_session_workspace_path(session_id).await; + let effective_path = self.effective_session_storage_path(session_id).await; let session_snapshot = self.sessions.get(session_id).map(|s| s.clone()); // Ref guard released -- DashMap shard lock is free. if let (Some(workspace_path), Some(session)) = (effective_path, session_snapshot) { @@ -2520,7 +2556,7 @@ impl SessionManager { session_id, elapsed_ms_u64(memory_stage_started_at) ); - self.session_workspace_index.remove(session_id); + self.session_storage_path_index.remove(session_id); info!( "Session deletion completed: session_id={}, workspace_path={}, duration_ms={}", @@ -2532,13 +2568,30 @@ impl SessionManager { Ok(()) } - /// Restore session (from persistent storage) + /// Restore session from a local or legacy workspace path. + /// + /// Callers that know remote identity must use [`Self::restore_session_for_workspace`]. + /// Callers that already resolved a `sessions` directory must use + /// [`Self::restore_session_from_storage_path`]. pub async fn restore_session( &self, workspace_path: &Path, session_id: &str, ) -> BitFunResult { - self.restore_session_internal(workspace_path, session_id, false) + let session_storage_path = self + .resolve_storage_path_for_workspace_path(workspace_path) + .await; + self.restore_session_from_storage_path(&session_storage_path, session_id) + .await + } + + pub async fn restore_session_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult { + let session_storage_path = self.resolve_storage_path_for_request(request).await?; + self.restore_session_from_storage_path(&session_storage_path, session_id) .await } @@ -2547,18 +2600,53 @@ impl SessionManager { workspace_path: &Path, session_id: &str, ) -> BitFunResult { - self.restore_session_internal(workspace_path, session_id, true) + let session_storage_path = self + .resolve_storage_path_for_workspace_path(workspace_path) + .await; + self.restore_internal_session_from_storage_path(&session_storage_path, session_id) .await } - async fn restore_session_internal( + pub async fn restore_internal_session_for_workspace( &self, - workspace_path: &Path, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult { + let session_storage_path = self.resolve_storage_path_for_request(request).await?; + self.restore_internal_session_from_storage_path(&session_storage_path, session_id) + .await + } + + pub async fn restore_session_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult { + self.restore_session_from_storage_path_internal(session_storage_path, session_id, false) + .await + } + + pub async fn restore_internal_session_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult { + self.restore_session_from_storage_path_internal(session_storage_path, session_id, true) + .await + } + + async fn restore_session_from_storage_path_internal( + &self, + session_storage_path: &Path, session_id: &str, include_internal: bool, ) -> BitFunResult { let (session, _) = self - .restore_session_with_turns_internal(workspace_path, session_id, include_internal) + .restore_session_with_turns_from_storage_path_internal( + session_storage_path, + session_id, + include_internal, + ) .await?; Ok(session) } @@ -2566,6 +2654,10 @@ impl SessionManager { /// Restore the persisted session header and turns needed by the UI view /// without loading runtime context snapshots or inserting the session into /// the in-memory coordinator state. + /// + /// This workspace-path overload is for local or legacy callers. Remote + /// callers must use [`Self::restore_session_view_for_workspace_timed`] or a + /// storage-path restore method so remote identity is preserved. pub async fn restore_session_view( &self, workspace_path: &Path, @@ -2581,9 +2673,31 @@ impl SessionManager { workspace_path: &Path, session_id: &str, ) -> BitFunResult<(Session, Vec, SessionViewRestoreTiming)> { - self.restore_session_view_internal(workspace_path, session_id, false, None) - .await - .map(|(session, turns, _, timing)| (session, turns, timing)) + let storage_path_started_at = Instant::now(); + let session_storage_path = self + .resolve_storage_path_for_workspace_path(workspace_path) + .await; + let resolve_storage_path_duration_ms = elapsed_ms_u64(storage_path_started_at); + let (session, turns, mut timing) = self + .restore_session_view_from_storage_path_timed(&session_storage_path, session_id) + .await?; + timing.resolve_storage_path_duration_ms = resolve_storage_path_duration_ms; + Ok((session, turns, timing)) + } + + pub async fn restore_session_view_for_workspace_timed( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult<(Session, Vec, SessionViewRestoreTiming)> { + let storage_path_started_at = Instant::now(); + let session_storage_path = self.resolve_storage_path_for_request(request).await?; + let resolve_storage_path_duration_ms = elapsed_ms_u64(storage_path_started_at); + let (session, turns, mut timing) = self + .restore_session_view_from_storage_path_timed(&session_storage_path, session_id) + .await?; + timing.resolve_storage_path_duration_ms = resolve_storage_path_duration_ms; + Ok((session, turns, timing)) } pub async fn restore_internal_session_view( @@ -2601,9 +2715,37 @@ impl SessionManager { workspace_path: &Path, session_id: &str, ) -> BitFunResult<(Session, Vec, SessionViewRestoreTiming)> { - self.restore_session_view_internal(workspace_path, session_id, true, None) - .await - .map(|(session, turns, _, timing)| (session, turns, timing)) + let storage_path_started_at = Instant::now(); + let session_storage_path = self + .resolve_storage_path_for_workspace_path(workspace_path) + .await; + let resolve_storage_path_duration_ms = elapsed_ms_u64(storage_path_started_at); + let (session, turns, mut timing) = self + .restore_internal_session_view_from_storage_path_timed( + &session_storage_path, + session_id, + ) + .await?; + timing.resolve_storage_path_duration_ms = resolve_storage_path_duration_ms; + Ok((session, turns, timing)) + } + + pub async fn restore_internal_session_view_for_workspace_timed( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult<(Session, Vec, SessionViewRestoreTiming)> { + let storage_path_started_at = Instant::now(); + let session_storage_path = self.resolve_storage_path_for_request(request).await?; + let resolve_storage_path_duration_ms = elapsed_ms_u64(storage_path_started_at); + let (session, turns, mut timing) = self + .restore_internal_session_view_from_storage_path_timed( + &session_storage_path, + session_id, + ) + .await?; + timing.resolve_storage_path_duration_ms = resolve_storage_path_duration_ms; + Ok((session, turns, timing)) } pub async fn restore_session_view_tail( @@ -2628,8 +2770,20 @@ impl SessionManager { usize, SessionViewRestoreTiming, )> { - self.restore_session_view_internal(workspace_path, session_id, false, Some(tail_turn_count)) - .await + let storage_path_started_at = Instant::now(); + let session_storage_path = self + .resolve_storage_path_for_workspace_path(workspace_path) + .await; + let resolve_storage_path_duration_ms = elapsed_ms_u64(storage_path_started_at); + let (session, turns, total_turn_count, mut timing) = self + .restore_session_view_from_storage_path_tail_timed( + &session_storage_path, + session_id, + tail_turn_count, + ) + .await?; + timing.resolve_storage_path_duration_ms = resolve_storage_path_duration_ms; + Ok((session, turns, total_turn_count, timing)) } pub async fn restore_internal_session_view_tail( @@ -2654,13 +2808,95 @@ impl SessionManager { usize, SessionViewRestoreTiming, )> { - self.restore_session_view_internal(workspace_path, session_id, true, Some(tail_turn_count)) - .await + let storage_path_started_at = Instant::now(); + let session_storage_path = self + .resolve_storage_path_for_workspace_path(workspace_path) + .await; + let resolve_storage_path_duration_ms = elapsed_ms_u64(storage_path_started_at); + let (session, turns, total_turn_count, mut timing) = self + .restore_internal_session_view_from_storage_path_tail_timed( + &session_storage_path, + session_id, + tail_turn_count, + ) + .await?; + timing.resolve_storage_path_duration_ms = resolve_storage_path_duration_ms; + Ok((session, turns, total_turn_count, timing)) } - async fn restore_session_view_internal( + pub async fn restore_session_view_from_storage_path_timed( &self, - workspace_path: &Path, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult<(Session, Vec, SessionViewRestoreTiming)> { + self.restore_session_view_from_storage_path_internal( + session_storage_path, + session_id, + false, + None, + ) + .await + .map(|(session, turns, _, timing)| (session, turns, timing)) + } + + pub async fn restore_internal_session_view_from_storage_path_timed( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult<(Session, Vec, SessionViewRestoreTiming)> { + self.restore_session_view_from_storage_path_internal( + session_storage_path, + session_id, + true, + None, + ) + .await + .map(|(session, turns, _, timing)| (session, turns, timing)) + } + + pub async fn restore_session_view_from_storage_path_tail_timed( + &self, + session_storage_path: &Path, + session_id: &str, + tail_turn_count: usize, + ) -> BitFunResult<( + Session, + Vec, + usize, + SessionViewRestoreTiming, + )> { + self.restore_session_view_from_storage_path_internal( + session_storage_path, + session_id, + false, + Some(tail_turn_count), + ) + .await + } + + pub async fn restore_internal_session_view_from_storage_path_tail_timed( + &self, + session_storage_path: &Path, + session_id: &str, + tail_turn_count: usize, + ) -> BitFunResult<( + Session, + Vec, + usize, + SessionViewRestoreTiming, + )> { + self.restore_session_view_from_storage_path_internal( + session_storage_path, + session_id, + true, + Some(tail_turn_count), + ) + .await + } + + async fn restore_session_view_from_storage_path_internal( + &self, + session_storage_path: &Path, session_id: &str, include_internal: bool, tail_turn_count: Option, @@ -2670,21 +2906,11 @@ impl SessionManager { usize, SessionViewRestoreTiming, )> { - let restore_request = SessionViewRestoreRequest { - workspace_path: workspace_path.to_path_buf(), - session_id: session_id.to_string(), - include_internal, - tail_turn_count, - }; let restore_started_at = Instant::now(); - let storage_path_started_at = Instant::now(); - let session_storage_path = self - .effective_storage_path_for_workspace_path(&restore_request.workspace_path) - .await; - let resolve_storage_path_duration_ms = elapsed_ms_u64(storage_path_started_at); + let resolve_storage_path_duration_ms = 0; debug!( - "Session view restore phase completed: session_id={}, phase=resolve_storage_path, duration_ms={}", - restore_request.session_id, + "Session view restore phase completed: session_id={}, phase=use_storage_path, duration_ms={}", + session_id, resolve_storage_path_duration_ms ); @@ -2693,39 +2919,34 @@ impl SessionManager { .persistence_manager .load_session_metadata(&session_storage_path, session_id) .await? - .is_some_and(|metadata| { - !restore_request.include_internal && metadata.should_hide_from_user_lists() - }) + .is_some_and(|metadata| !include_internal && metadata.should_hide_from_user_lists()) { return Err(BitFunError::NotFound(format!( "Session not found: {}", - restore_request.session_id + session_id ))); } let visibility_metadata_duration_ms = elapsed_ms_u64(metadata_started_at); debug!( "Session view restore phase completed: session_id={}, phase=load_metadata, duration_ms={}", - restore_request.session_id, + session_id, visibility_metadata_duration_ms ); let session_started_at = Instant::now(); let (mut session, persisted_turns, total_turn_count, turn_load) = - if let Some(tail_turn_count) = restore_request.tail_turn_count { + if let Some(tail_turn_count) = tail_turn_count { self.persistence_manager .load_session_with_tail_turns_timed( &session_storage_path, - &restore_request.session_id, + session_id, tail_turn_count, ) .await? } else { let (session, turns, timing) = self .persistence_manager - .load_session_with_turns_timed( - &session_storage_path, - &restore_request.session_id, - ) + .load_session_with_turns_timed(&session_storage_path, session_id) .await?; let total_turn_count = turns.len(); (session, turns, total_turn_count, timing) @@ -2736,7 +2957,7 @@ impl SessionManager { session_id, persisted_turns.len(), total_turn_count, - restore_request.tail_turn_count, + tail_turn_count, load_session_with_turns_duration_ms ); @@ -2787,42 +3008,99 @@ impl SessionManager { } /// Restore session and return the persisted turns read during restore. + /// + /// This workspace-path overload is for local or legacy callers. Remote + /// callers must use [`Self::restore_session_with_turns_for_workspace`] or a + /// storage-path restore method so remote identity is preserved. pub async fn restore_session_with_turns( &self, workspace_path: &Path, session_id: &str, ) -> BitFunResult<(Session, Vec)> { - self.restore_session_with_turns_internal(workspace_path, session_id, false) + let session_storage_path = self + .resolve_storage_path_for_workspace_path(workspace_path) + .await; + self.restore_session_with_turns_from_storage_path(&session_storage_path, session_id) .await } - pub async fn restore_internal_session_with_turns( + pub async fn restore_session_with_turns_for_workspace( &self, - workspace_path: &Path, + request: SessionStoragePathRequest, session_id: &str, ) -> BitFunResult<(Session, Vec)> { - self.restore_session_with_turns_internal(workspace_path, session_id, true) + let session_storage_path = self.resolve_storage_path_for_request(request).await?; + self.restore_session_with_turns_from_storage_path(&session_storage_path, session_id) .await } - async fn restore_session_with_turns_internal( + pub async fn restore_internal_session_with_turns( &self, workspace_path: &Path, session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + let session_storage_path = self + .resolve_storage_path_for_workspace_path(workspace_path) + .await; + self.restore_internal_session_with_turns_from_storage_path( + &session_storage_path, + session_id, + ) + .await + } + + pub async fn restore_internal_session_with_turns_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + let session_storage_path = self.resolve_storage_path_for_request(request).await?; + self.restore_internal_session_with_turns_from_storage_path( + &session_storage_path, + session_id, + ) + .await + } + + pub async fn restore_session_with_turns_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + self.restore_session_with_turns_from_storage_path_internal( + session_storage_path, + session_id, + false, + ) + .await + } + + pub async fn restore_internal_session_with_turns_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + self.restore_session_with_turns_from_storage_path_internal( + session_storage_path, + session_id, + true, + ) + .await + } + + async fn restore_session_with_turns_from_storage_path_internal( + &self, + session_storage_path: &Path, + session_id: &str, include_internal: bool, ) -> BitFunResult<(Session, Vec)> { let restore_started_at = Instant::now(); // Check if session is already in memory let session_already_in_memory = self.sessions.contains_key(session_id); - let storage_path_started_at = Instant::now(); - let session_storage_path = self - .effective_storage_path_for_workspace_path(workspace_path) - .await; debug!( - "Session restore phase completed: session_id={}, phase=resolve_storage_path, duration_ms={}", - session_id, - elapsed_ms_u64(storage_path_started_at) + "Session restore phase completed: session_id={}, phase=use_storage_path, duration_ms=0", + session_id ); let metadata_started_at = Instant::now(); @@ -3079,8 +3357,8 @@ impl SessionManager { // 4. Add to memory (will overwrite if already exists) self.sessions .insert(session_id.to_string(), session.clone()); - self.session_workspace_index - .insert(session_id.to_string(), session_storage_path.clone()); + self.session_storage_path_index + .insert(session_id.to_string(), session_storage_path.to_path_buf()); Ok((session, persisted_turns)) } @@ -3275,7 +3553,7 @@ impl SessionManager { ))); } - self.effective_session_workspace_path(session_id) + self.effective_session_storage_path(session_id) .await .ok_or_else(|| { BitFunError::Validation(format!( @@ -3762,7 +4040,7 @@ impl SessionManager { } let workspace_path = self - .effective_session_workspace_path(session_id) + .effective_session_storage_path(session_id) .await .ok_or_else(|| { BitFunError::Validation(format!( @@ -3874,7 +4152,7 @@ impl SessionManager { } let workspace_path = self - .effective_session_workspace_path(session_id) + .effective_session_storage_path(session_id) .await .ok_or_else(|| { BitFunError::Validation(format!( @@ -3936,7 +4214,7 @@ impl SessionManager { } let workspace_path = self - .effective_session_workspace_path(session_id) + .effective_session_storage_path(session_id) .await .ok_or_else(|| { BitFunError::Validation(format!( @@ -4002,7 +4280,7 @@ impl SessionManager { } let workspace_path = self - .effective_session_workspace_path(session_id) + .effective_session_storage_path(session_id) .await .ok_or_else(|| { BitFunError::Validation(format!( @@ -4066,7 +4344,7 @@ impl SessionManager { } let workspace_path = self - .effective_session_workspace_path(session_id) + .effective_session_storage_path(session_id) .await .ok_or_else(|| { BitFunError::Validation(format!( @@ -4122,7 +4400,7 @@ impl SessionManager { /// canonical turn history instead of the runtime context cache. pub async fn get_messages(&self, session_id: &str) -> BitFunResult> { if self.config.enable_persistence { - if let Some(workspace_path) = self.effective_session_workspace_path(session_id).await { + if let Some(workspace_path) = self.effective_session_storage_path(session_id).await { let messages = self .rebuild_messages_from_turns(&workspace_path, session_id) .await?; @@ -4205,7 +4483,7 @@ impl SessionManager { session_id: &str, compression_state: CompressionState, ) -> BitFunResult<()> { - let effective_path = self.effective_session_workspace_path(session_id).await; + let effective_path = self.effective_session_storage_path(session_id).await; // IMPORTANT: keep the DashMap guard scope short -- do NOT hold it across .await. let session_snapshot = if let Some(mut session) = self.sessions.get_mut(session_id) { @@ -4522,6 +4800,7 @@ mod tests { SessionRelationship, SessionRelationshipKind, ToolCallData, ToolItemData, ToolResultData, TurnStatus, UserMessageData, }; + use bitfun_runtime_ports::SessionStoragePathRequest; use dashmap::try_result::TryResult; use serde_json::json; use std::collections::HashSet; @@ -5283,6 +5562,113 @@ mod tests { ); } + #[tokio::test] + async fn restore_session_from_storage_path_accepts_resolved_sessions_dir() { + let workspace = TestWorkspace::new(); + let path_manager = workspace.path_manager(); + let persistence_manager = + Arc::new(PersistenceManager::new(path_manager.clone()).expect("persistence manager")); + let manager = test_manager(persistence_manager.clone()); + let sessions_dir = path_manager.project_sessions_dir(workspace.path()); + let session_id = Uuid::new_v4().to_string(); + let session = Session::new_with_id( + session_id.clone(), + "Resolved sessions restore".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ); + + persistence_manager + .save_session(&sessions_dir, &session) + .await + .expect("session should save to resolved sessions dir"); + + let restored = manager + .restore_session_from_storage_path(&sessions_dir, &session_id) + .await + .expect("storage restore should read the resolved sessions dir directly"); + + assert_eq!(restored.session_id, session_id); + } + + #[tokio::test] + async fn restore_session_workspace_api_does_not_accept_resolved_sessions_dir() { + let workspace = TestWorkspace::new(); + let path_manager = workspace.path_manager(); + let persistence_manager = + Arc::new(PersistenceManager::new(path_manager.clone()).expect("persistence manager")); + let manager = test_manager(persistence_manager.clone()); + let sessions_dir = path_manager.project_sessions_dir(workspace.path()); + let session_id = Uuid::new_v4().to_string(); + let session = Session::new_with_id( + session_id.clone(), + "Resolved sessions restore".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ); + + persistence_manager + .save_session(&sessions_dir, &session) + .await + .expect("session should save to resolved sessions dir"); + + let result = manager.restore_session(&sessions_dir, &session_id).await; + + assert!( + result.is_err(), + "workspace restore should not accept an already-resolved sessions dir" + ); + } + + #[tokio::test] + async fn restore_session_for_workspace_uses_remote_identity() { + let workspace = TestWorkspace::new(); + let path_manager = workspace.path_manager(); + let persistence_manager = + Arc::new(PersistenceManager::new(path_manager.clone()).expect("persistence manager")); + let manager = test_manager(persistence_manager.clone()); + let sessions_dir = crate::service::WorkspaceRuntimeService::new(path_manager.clone()) + .context_for_remote_workspace("dev-host", "/home/wsp/project") + .sessions_dir; + let session_id = Uuid::new_v4().to_string(); + let session = Session::new_with_id( + session_id.clone(), + "Remote identity restore".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some("/home/wsp/project".to_string()), + remote_connection_id: Some("ssh-1".to_string()), + remote_ssh_host: Some("dev-host".to_string()), + ..Default::default() + }, + ); + + persistence_manager + .save_session(&sessions_dir, &session) + .await + .expect("session should save to remote sessions dir"); + + let restored = manager + .restore_session_for_workspace( + SessionStoragePathRequest { + workspace_path: PathBuf::from("/home/wsp/project"), + remote_connection_id: Some("ssh-1".to_string()), + remote_ssh_host: Some("dev-host".to_string()), + }, + &session_id, + ) + .await + .expect("workspace restore should use remote identity"); + + assert_eq!(restored.session_id, session_id); + } + #[tokio::test] async fn restore_session_view_loads_turns_without_restoring_runtime_context() { let workspace = TestWorkspace::new(); @@ -6189,7 +6575,7 @@ mod tests { assert_eq!( manager - .session_workspace_index + .session_storage_path_index .get(&session.session_id) .as_deref() .map(|entry| entry.to_path_buf()), @@ -6202,7 +6588,7 @@ mod tests { .expect("session should delete"); assert!(manager - .session_workspace_index + .session_storage_path_index .get(&session.session_id) .is_none()); } diff --git a/src/crates/assembly/core/src/agentic/session/session_store_port.rs b/src/crates/assembly/core/src/agentic/session/session_store_port.rs index 607994b6c..b1466f36a 100644 --- a/src/crates/assembly/core/src/agentic/session/session_store_port.rs +++ b/src/crates/assembly/core/src/agentic/session/session_store_port.rs @@ -20,13 +20,17 @@ pub struct CoreSessionStorePort { } impl CoreSessionStorePort { - #[cfg(test)] - pub fn with_path_manager_for_tests(path_manager: Arc) -> Self { + pub(crate) fn with_path_manager(path_manager: Arc) -> Self { Self { path_manager: Some(path_manager), } } + #[cfg(test)] + pub fn with_path_manager_for_tests(path_manager: Arc) -> Self { + Self::with_path_manager(path_manager) + } + fn path_manager(&self) -> Arc { self.path_manager .clone() diff --git a/src/crates/assembly/core/src/service/remote_connect/remote_server.rs b/src/crates/assembly/core/src/service/remote_connect/remote_server.rs index f11d56acf..df591e87d 100644 --- a/src/crates/assembly/core/src/service/remote_connect/remote_server.rs +++ b/src/crates/assembly/core/src/service/remote_connect/remote_server.rs @@ -469,7 +469,7 @@ mod tests { assert_eq!( remote_session_restore_target(false, Some(&binding)), - Some("/workspace/project") + Some(binding.clone()) ); assert_eq!(remote_session_restore_target(true, Some(&binding)), None); assert_eq!(remote_session_restore_target(false, None), None); diff --git a/src/crates/assembly/core/src/service_agent_runtime.rs b/src/crates/assembly/core/src/service_agent_runtime.rs index 05a3adf53..780f2a04e 100644 --- a/src/crates/assembly/core/src/service_agent_runtime.rs +++ b/src/crates/assembly/core/src/service_agent_runtime.rs @@ -399,7 +399,7 @@ async fn resolve_session_model_id(session_id: &str) -> Option { let session_storage_dir = CoreServiceAgentRuntime::resolve_session_storage_dir(session_id).await?; coordinator - .restore_session(&session_storage_dir, session_id) + .restore_session_from_storage_path(&session_storage_dir, session_id) .await .ok() .and_then(|session| normalize_remote_session_model_id(session.config.model_id.clone())) @@ -686,7 +686,7 @@ impl CoreServiceAgentRuntime { )); }; coordinator - .restore_session(&session_storage_dir, session_id) + .restore_session_from_storage_path(&session_storage_dir, session_id) .await .map_err(|e| format!("Failed to restore session: {e}"))?; } @@ -1022,16 +1022,28 @@ impl RemoteDialogRuntimeHost for CoreRemoteDialogRuntimeHost<'_> { async fn restore_remote_session( &self, session_id: &str, - workspace_path: &str, + workspace: RemoteDialogWorkspaceBinding, ) -> Result<(), String> { - let restore_path = CoreServiceAgentRuntime::resolve_session_storage_dir(session_id) - .await - .unwrap_or_else(|| std::path::PathBuf::from(workspace_path)); - self.coordinator - .restore_session(&restore_path, session_id) - .await - .map(|_| ()) - .map_err(|e| e.to_string()) + if let Some(session_storage_dir) = + CoreServiceAgentRuntime::resolve_session_storage_dir(session_id).await + { + self.coordinator + .restore_session_from_storage_path(&session_storage_dir, session_id) + .await + } else { + self.coordinator + .restore_session_for_workspace( + SessionStoragePathRequest { + workspace_path: std::path::PathBuf::from(workspace.workspace_path), + remote_connection_id: workspace.remote_connection_id, + remote_ssh_host: workspace.remote_ssh_host, + }, + session_id, + ) + .await + } + .map(|_| ()) + .map_err(|e| e.to_string()) } fn prewarm_remote_terminal(&self, request: RemoteTerminalPrewarmRequest) { @@ -1271,7 +1283,7 @@ impl RemoteSessionRuntimeHost for CoreRemoteSessionRuntimeHost { )); }; self.coordinator - .restore_session(&session_storage_dir, session_id) + .restore_session_from_storage_path(&session_storage_dir, session_id) .await .map(|_| ()) .map_err(|error| format!("Failed to restore session: {error}")) @@ -1406,7 +1418,7 @@ impl RemoteCancelRuntimeHost for CoreRemoteCancelRuntimeHost { .await .unwrap_or_else(|| std::path::PathBuf::from(restore_path_hint)); self.coordinator - .restore_session(&restore_path, session_id) + .restore_session_from_storage_path(&restore_path, session_id) .await .map(|_| ()) .map_err(|error| error.to_string()) diff --git a/src/crates/services/services-integrations/src/remote_connect.rs b/src/crates/services/services-integrations/src/remote_connect.rs index e6fd0b3d6..6e5af069d 100644 --- a/src/crates/services/services-integrations/src/remote_connect.rs +++ b/src/crates/services/services-integrations/src/remote_connect.rs @@ -189,14 +189,14 @@ pub fn resolve_remote_execution_image_contexts( image_contexts.unwrap_or_else(|| legacy_contexts(legacy_images)) } -pub fn remote_session_restore_target<'a>( +pub fn remote_session_restore_target( session_exists: bool, - binding_workspace: Option<&'a RemoteDialogWorkspaceBinding>, -) -> Option<&'a str> { + binding_workspace: Option<&RemoteDialogWorkspaceBinding>, +) -> Option { if session_exists { None } else { - binding_workspace.map(|binding| binding.workspace_path.as_str()) + binding_workspace.cloned() } } @@ -412,7 +412,7 @@ pub trait RemoteDialogRuntimeHost: Send + Sync { async fn restore_remote_session( &self, session_id: &str, - workspace_path: &str, + workspace: RemoteDialogWorkspaceBinding, ) -> Result<(), String>; fn prewarm_remote_terminal(&self, request: RemoteTerminalPrewarmRequest); @@ -446,12 +446,10 @@ where let binding_workspace = host.resolve_binding_workspace(&session_id).await; let session_exists = host.remote_session_exists(&session_id).await?; - if let Some(workspace_path) = + if let Some(workspace) = remote_session_restore_target(session_exists, binding_workspace.as_ref()) { - let _ = host - .restore_remote_session(&session_id, workspace_path) - .await; + let _ = host.restore_remote_session(&session_id, workspace).await; } host.prewarm_remote_terminal(RemoteTerminalPrewarmRequest { diff --git a/src/crates/services/services-integrations/tests/remote_connect_contracts.rs b/src/crates/services/services-integrations/tests/remote_connect_contracts.rs index 9b5104f4d..d67cdeac4 100644 --- a/src/crates/services/services-integrations/tests/remote_connect_contracts.rs +++ b/src/crates/services/services-integrations/tests/remote_connect_contracts.rs @@ -335,7 +335,7 @@ fn remote_connect_cancel_and_restore_policy_preserve_runtime_decisions() { let binding = RemoteDialogWorkspaceBinding::local("D:/workspace/project"); assert_eq!( remote_session_restore_target(false, Some(&binding)), - Some("D:/workspace/project") + Some(binding.clone()) ); assert_eq!(remote_session_restore_target(true, Some(&binding)), None); assert_eq!(remote_session_restore_target(false, None), None); @@ -508,12 +508,18 @@ impl RemoteDialogRuntimeHost for RecordingDialogHost { async fn restore_remote_session( &self, session_id: &str, - workspace_path: &str, + workspace: RemoteDialogWorkspaceBinding, ) -> Result<(), String> { - self.events - .lock() - .unwrap() - .push(format!("restore:{session_id}:{workspace_path}")); + self.events.lock().unwrap().push(format!( + "restore:{}:{}:{}:{}", + session_id, + workspace.workspace_path, + workspace + .remote_connection_id + .as_deref() + .unwrap_or(""), + workspace.remote_ssh_host.as_deref().unwrap_or("") + )); if self.restore_error { Err("restore failed".to_string()) } else { @@ -945,7 +951,7 @@ async fn remote_connect_dialog_runtime_owns_restore_prewarm_and_submit_order() { "ensure_tracker:session-1", "resolve_workspace:session-1", "session_exists:session-1", - "restore:session-1:D:/workspace/project", + "restore:session-1:D:/workspace/project::", "prewarm:session-1:D:/workspace/project", "generate_turn", "submit:session-1", @@ -998,6 +1004,18 @@ async fn remote_connect_dialog_runtime_preserves_remote_workspace_identity() { .await .expect("dialog submit succeeds"); + assert_eq!( + host.events(), + vec![ + "ensure_tracker:session-1", + "resolve_workspace:session-1", + "session_exists:session-1", + "restore:session-1:/home/wsp/project:ssh-1:dev-host", + "prewarm:session-1:/home/wsp/project", + "submit:session-1", + ] + ); + let submitted = host.submitted(); let binding = submitted .binding_workspace @@ -1103,7 +1121,7 @@ async fn remote_connect_dialog_runtime_keeps_legacy_restore_failure_tolerance() "ensure_tracker:session-1", "resolve_workspace:session-1", "session_exists:session-1", - "restore:session-1:D:/workspace/project", + "restore:session-1:D:/workspace/project::", "prewarm:session-1:D:/workspace/project", "submit:session-1", ] From 3a714ef15b8e6c7cf479d9e2cebef9d4d10416e0 Mon Sep 17 00:00:00 2001 From: wsp Date: Fri, 26 Jun 2026 22:31:41 +0800 Subject: [PATCH 4/4] fix: shorten generated cron job ids Generate scheduled job ids as cron_ plus an 8-character lowercase hex suffix instead of a full UUID string. Keep creation collision-safe by checking existing jobs and retrying when a generated id is already present. Add coverage for the shortened id format. --- .../assembly/core/src/service/cron/service.rs | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/src/crates/assembly/core/src/service/cron/service.rs b/src/crates/assembly/core/src/service/cron/service.rs index 78ffb3ed1..49aaf3777 100644 --- a/src/crates/assembly/core/src/service/cron/service.rs +++ b/src/crates/assembly/core/src/service/cron/service.rs @@ -155,7 +155,7 @@ impl CronService { validate_schedule(&schedule, current_ms)?; let mut job = CronJob { - id: format!("cron_{}", Uuid::new_v4().simple()), + id: generate_cron_job_id(&jobs), name: request.name.trim().to_string(), schedule, payload: request.payload, @@ -963,6 +963,16 @@ fn now_ms() -> i64 { Utc::now().timestamp_millis() } +fn generate_cron_job_id(jobs: &HashMap) -> String { + loop { + let uuid = Uuid::new_v4().simple().to_string(); + let id = format!("cron_{}", &uuid[..8]); + if !jobs.contains_key(&id) { + return id; + } + } +} + struct EnqueueInput { job_id: String, job_name: String, @@ -996,6 +1006,18 @@ fn cron_enqueue_error_is_missing_session(error: &str) -> bool { mod tests { use super::*; + #[test] + fn generate_cron_job_id_uses_short_hex_suffix() { + let jobs = HashMap::new(); + let id = generate_cron_job_id(&jobs); + + assert_eq!(id.len(), "cron_".len() + 8); + assert!(id.starts_with("cron_")); + assert!(id["cron_".len()..] + .chars() + .all(|ch| ch.is_ascii_hexdigit() && !ch.is_ascii_uppercase())); + } + #[test] fn materialize_workspace_ref_normalizes_windows_style_paths() { let workspace = materialize_workspace_ref(CronWorkspaceRef {