Skip to content

Commit 935a9a4

Browse files
authored
feat(desktop,core): remote workspace tool context and snapshot/stream fixes (#602)
- Wire desktop tool API ToolUseContext with remote WorkspaceBinding and workspace services - Propagate workspace_services into execution engine tool availability context - Ignore whitespace-only stream chunks for effective output / first visible timing - Short-circuit snapshot session files and stats for remote workspace paths - Add unit tests for remote skill discovery and whitespace stream handling
1 parent 610894f commit 935a9a4

5 files changed

Lines changed: 259 additions & 25 deletions

File tree

src/apps/desktop/src/api/snapshot_service.rs

Lines changed: 24 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@
33
use bitfun_core::infrastructure::try_get_path_manager_arc;
44
use bitfun_core::service::remote_ssh::workspace_state::is_remote_path;
55
use bitfun_core::service::snapshot::{
6-
ensure_snapshot_manager_for_workspace, get_snapshot_manager_for_workspace,
7-
initialize_snapshot_manager_for_workspace, OperationType, SnapshotConfig, SnapshotManager,
6+
OperationType, SnapshotConfig, SnapshotManager, ensure_snapshot_manager_for_workspace,
7+
get_snapshot_manager_for_workspace, initialize_snapshot_manager_for_workspace,
88
};
99
use log::{info, warn};
1010
use serde::{Deserialize, Serialize};
@@ -285,7 +285,7 @@ pub async fn record_file_change(
285285
return Err(format!(
286286
"Unknown operation type: {}",
287287
request.operation_type
288-
))
288+
));
289289
}
290290
};
291291

@@ -424,7 +424,10 @@ pub async fn rollback_to_turn(
424424
deleted_turns_count = count;
425425
}
426426
Err(e) => {
427-
warn!("Failed to delete conversation turns: session_id={}, turn_index={}, error={}", request.session_id, request.turn_index, e);
427+
warn!(
428+
"Failed to delete conversation turns: session_id={}, turn_index={}, error={}",
429+
request.session_id, request.turn_index, e
430+
);
428431
}
429432
}
430433
}
@@ -546,6 +549,10 @@ pub async fn reject_file(
546549

547550
#[tauri::command]
548551
pub async fn get_session_files(request: GetSessionFilesRequest) -> Result<Vec<String>, String> {
552+
if is_remote_path(&request.workspace_path).await {
553+
return Ok(vec![]);
554+
}
555+
549556
let manager = ensure_snapshot_manager_ready(&request.workspace_path).await?;
550557

551558
let files = manager
@@ -580,7 +587,10 @@ pub async fn get_session_turns(
580587
}
581588
Ok(None) => {}
582589
Err(e) => {
583-
warn!("Failed to load conversation metadata: session_id={}, error={}, falling back to snapshot", request.session_id, e);
590+
warn!(
591+
"Failed to load conversation metadata: session_id={}, error={}, falling back to snapshot",
592+
request.session_id, e
593+
);
584594
}
585595
}
586596
}
@@ -840,6 +850,15 @@ pub async fn reject_operation(
840850
pub async fn get_session_stats(
841851
request: GetSessionStatsRequest,
842852
) -> Result<serde_json::Value, String> {
853+
if is_remote_path(&request.workspace_path).await {
854+
return Ok(serde_json::json!({
855+
"session_id": request.session_id,
856+
"total_files": 0,
857+
"total_turns": 0,
858+
"total_changes": 0
859+
}));
860+
}
861+
843862
let manager = ensure_snapshot_manager_ready(&request.workspace_path).await?;
844863

845864
let stats = manager

src/apps/desktop/src/api/tool_api.rs

Lines changed: 59 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,13 @@ use std::collections::HashMap;
66
use std::path::PathBuf;
77

88
use bitfun_core::agentic::{
9+
WorkspaceBinding,
910
tools::framework::ToolUseContext,
1011
tools::{get_all_tools, get_readonly_tools},
11-
WorkspaceBinding,
12+
workspace::{local_workspace_services, remote_workspace_services},
13+
};
14+
use bitfun_core::service::remote_ssh::workspace_state::{
15+
get_remote_workspace_manager, lookup_remote_connection, resolve_workspace_session_identity,
1216
};
1317
use bitfun_core::util::elapsed_ms_u64;
1418

@@ -82,23 +86,71 @@ pub struct ToolConfirmationResponse {
8286
pub message: String,
8387
}
8488

85-
fn build_tool_context(workspace_path: Option<&str>) -> ToolUseContext {
89+
async fn build_tool_context(workspace_path: Option<&str>) -> ToolUseContext {
8690
let normalized_workspace_path = workspace_path
8791
.map(str::trim)
8892
.filter(|path| !path.is_empty());
8993

94+
let workspace = match normalized_workspace_path {
95+
Some(path) => match resolve_workspace_session_identity(path, None, None).await {
96+
Some(identity) => {
97+
if let Some(connection_id) = identity.remote_connection_id.as_deref() {
98+
let connection_name = lookup_remote_connection(path)
99+
.await
100+
.map(|entry| entry.connection_name)
101+
.unwrap_or_else(|| connection_id.to_string());
102+
Some(WorkspaceBinding::new_remote(
103+
None,
104+
PathBuf::from(path),
105+
connection_id.to_string(),
106+
connection_name,
107+
identity,
108+
))
109+
} else {
110+
Some(WorkspaceBinding::new(None, PathBuf::from(path)))
111+
}
112+
}
113+
None => Some(WorkspaceBinding::new(None, PathBuf::from(path))),
114+
},
115+
None => None,
116+
};
117+
118+
let workspace_services = match workspace.as_ref() {
119+
Some(binding) if binding.is_remote() => {
120+
let connection_id = binding.connection_id().map(str::to_string);
121+
match (connection_id, get_remote_workspace_manager()) {
122+
(Some(connection_id), Some(manager)) => {
123+
match (
124+
manager.get_file_service().await,
125+
manager.get_ssh_manager().await,
126+
) {
127+
(Some(file_service), Some(ssh_manager)) => Some(remote_workspace_services(
128+
connection_id,
129+
file_service,
130+
ssh_manager,
131+
binding.root_path_string(),
132+
)),
133+
_ => None,
134+
}
135+
}
136+
_ => None,
137+
}
138+
}
139+
Some(binding) => Some(local_workspace_services(binding.root_path_string())),
140+
None => None,
141+
};
142+
90143
ToolUseContext {
91144
tool_call_id: None,
92145
agent_type: None,
93146
session_id: None,
94147
dialog_turn_id: None,
95-
workspace: normalized_workspace_path
96-
.map(|path| WorkspaceBinding::new(None, PathBuf::from(path))),
148+
workspace,
97149
custom_data: HashMap::new(),
98150
computer_use_host: None,
99151
cancellation_token: None,
100152
runtime_tool_restrictions: Default::default(),
101-
workspace_services: None,
153+
workspace_services,
102154
}
103155
}
104156

@@ -229,7 +281,7 @@ pub async fn validate_tool_input(
229281
request.workspace_path.as_deref(),
230282
)?;
231283

232-
let context = build_tool_context(request.workspace_path.as_deref());
284+
let context = build_tool_context(request.workspace_path.as_deref()).await;
233285

234286
let validation_result = tool.validate_input(&request.input, Some(&context)).await;
235287

@@ -260,7 +312,7 @@ pub async fn execute_tool(request: ToolExecutionRequest) -> Result<ToolExecution
260312
request.workspace_path.as_deref(),
261313
)?;
262314

263-
let context = build_tool_context(request.workspace_path.as_deref());
315+
let context = build_tool_context(request.workspace_path.as_deref()).await;
264316

265317
let validation_result = tool.validate_input(&request.input, Some(&context)).await;
266318
if !validation_result.result {

src/crates/core/src/agentic/execution/execution_engine.rs

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,21 +5,21 @@
55
use super::round_executor::RoundExecutor;
66
use super::types::{ExecutionContext, ExecutionResult, RoundContext};
77
use crate::agentic::agents::{
8-
get_agent_registry, PromptBuilder, PromptBuilderContext, RemoteExecutionHints,
8+
PromptBuilder, PromptBuilderContext, RemoteExecutionHints, get_agent_registry,
99
};
1010
use crate::agentic::core::{
11-
render_system_reminder, Message, MessageContent, MessageHelper, MessageRole,
12-
MessageSemanticKind, RequestReasoningTokenPolicy, Session,
11+
Message, MessageContent, MessageHelper, MessageRole, MessageSemanticKind,
12+
RequestReasoningTokenPolicy, Session, render_system_reminder,
1313
};
1414
use crate::agentic::events::{AgenticEvent, EventPriority, EventQueue};
1515
use crate::agentic::execution::types::FinishReason;
1616
use crate::agentic::image_analysis::{
17-
build_multimodal_message_with_images, process_image_contexts_for_provider, ImageContextData,
18-
ImageLimits,
17+
ImageContextData, ImageLimits, build_multimodal_message_with_images,
18+
process_image_contexts_for_provider,
1919
};
2020
use crate::agentic::session::{CompressionTailPolicy, ContextCompressor, SessionManager};
2121
use crate::agentic::tools::{
22-
get_all_registered_tools, SubagentParentInfo, ToolRuntimeRestrictions,
22+
SubagentParentInfo, ToolRuntimeRestrictions, get_all_registered_tools,
2323
};
2424
use crate::agentic::util::build_remote_workspace_layout_preview;
2525
use crate::agentic::{WorkspaceBackend, WorkspaceBinding};
@@ -1272,6 +1272,7 @@ impl ExecutionEngine {
12721272
self.get_available_tools_and_definitions(
12731273
&allowed_tools,
12741274
context.workspace.as_ref(),
1275+
context.workspace_services.as_ref(),
12751276
&agent_type,
12761277
primary_supports_image_understanding,
12771278
)
@@ -1870,6 +1871,7 @@ impl ExecutionEngine {
18701871
&self,
18711872
mode_allowed_tools: &[String],
18721873
workspace: Option<&crate::agentic::WorkspaceBinding>,
1874+
workspace_services: Option<&crate::agentic::workspace::WorkspaceServices>,
18731875
agent_type: &str,
18741876
primary_supports_image_understanding: bool,
18751877
) -> (Vec<String>, Option<Vec<ToolDefinition>>) {
@@ -1893,7 +1895,7 @@ impl ExecutionEngine {
18931895
computer_use_host: None,
18941896
cancellation_token: None,
18951897
runtime_tool_restrictions: ToolRuntimeRestrictions::default(),
1896-
workspace_services: None,
1898+
workspace_services: workspace_services.cloned(),
18971899
};
18981900
for tool in &all_tools {
18991901
if !tool.is_enabled().await {

src/crates/core/src/agentic/execution/stream_processor.rs

Lines changed: 33 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -430,8 +430,7 @@ impl StreamProcessor {
430430
for tool_call in tool_calls {
431431
trace!(
432432
"Cleaning up tool: {} ({})",
433-
tool_call.tool_name,
434-
tool_call.tool_id
433+
tool_call.tool_name, tool_call.tool_id
435434
);
436435

437436
let tool_event = if is_user_cancellation {
@@ -583,8 +582,10 @@ impl StreamProcessor {
583582

584583
/// Handle text chunk
585584
async fn handle_text_chunk(&self, ctx: &mut StreamContext, text: String) {
586-
ctx.has_effective_output = true;
587-
ctx.mark_first_visible_output();
585+
if !text.trim().is_empty() {
586+
ctx.has_effective_output = true;
587+
ctx.mark_first_visible_output();
588+
}
588589
ctx.full_text.push_str(&text);
589590
ctx.text_chunks_count += 1;
590591

@@ -976,6 +977,34 @@ mod tests {
976977
assert_eq!(result.usage.as_ref().map(|u| u.total_token_count), Some(7));
977978
}
978979

980+
#[tokio::test]
981+
async fn whitespace_only_text_is_not_effective_output() {
982+
let processor = build_processor();
983+
let stream = iter(vec![Ok(UnifiedResponse {
984+
text: Some("\n\n ".to_string()),
985+
..Default::default()
986+
})])
987+
.boxed();
988+
989+
let result = processor
990+
.process_stream(
991+
stream,
992+
None,
993+
None,
994+
"session_1".to_string(),
995+
"turn_1".to_string(),
996+
"round_1".to_string(),
997+
None,
998+
&CancellationToken::new(),
999+
)
1000+
.await
1001+
.expect("stream result");
1002+
1003+
assert_eq!(result.full_text, "\n\n ");
1004+
assert!(!result.has_effective_output);
1005+
assert_eq!(result.first_visible_output_ms, None);
1006+
}
1007+
9791008
#[tokio::test]
9801009
async fn finalizes_tool_after_same_chunk_finish_reason() {
9811010
let processor = build_processor();

0 commit comments

Comments
 (0)