Skip to content

feat(hooks): support multiple client-scoped commands #758

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
1jehuang wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat(hooks): support multiple client-scoped commands #758

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
277 changes: 160 additions & 117 deletions crates/jcode-app-core/src/server/client_lifecycle.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,13 @@ fn initial_subscribe_working_dir(request: &Request) -> std::result::Result<Strin
}
}

fn subscribe_terminal_env(request: &Request) -> Vec<(String, String)> {
match request {
Request::Subscribe { terminal_env, .. } => terminal_env.clone(),
_ => Vec::new(),
}
}

struct ProcessingMessage {
id: u64,
content: String,
Expand Down Expand Up @@ -468,6 +475,7 @@ pub(super) async fn handle_client(
return Ok(());
}
};
let mut active_terminal_env = subscribe_terminal_env(&initial_request);

// Per-client state
let mut client_is_processing = false;
Expand All @@ -493,11 +501,15 @@ pub(super) async fn handle_client(

// Create a new session for this client
let t0 = std::time::Instant::now();
let mut new_agent = Agent::new_with_initial_working_dir(
Arc::clone(&provider),
registry.clone(),
Some(&initial_working_dir),
);
let mut new_agent =
crate::hooks::with_client_terminal_env(active_terminal_env.clone(), async {
Agent::new_with_initial_working_dir(
Arc::clone(&provider),
registry.clone(),
Some(&initial_working_dir),
)
})
.await;
let agent_new_ms = t0.elapsed().as_millis();

new_agent.set_memory_enabled(crate::config::config().features.memory);
Expand Down Expand Up @@ -525,7 +537,7 @@ pub(super) async fn handle_client(
last_seen: connected_at,
is_processing: false,
current_tool_name: None,
terminal_env: Vec::new(),
terminal_env: active_terminal_env.clone(),
disconnect_tx: disconnect_tx.clone(),
},
);
Expand Down Expand Up @@ -1132,6 +1144,7 @@ pub(super) async fn handle_client(
&agent,
&client_event_tx,
&processing_done_tx,
active_terminal_env.clone(),
&SwarmStatusRefs {
members: &swarm_members,
swarms_by_id: &swarms_by_id,
Expand Down Expand Up @@ -1209,28 +1222,31 @@ pub(super) async fn handle_client(
) {
continue;
}
handle_clear_session(
id,
client_selfdev,
&mut client_session_id,
&client_connection_id,
&agent,
&provider,
&registry,
&sessions,
&shutdown_signals,
&soft_interrupt_queues,
&client_connections,
&swarm_members,
&swarms_by_id,
&file_touch,
&channel_subscriptions,
&channel_subscriptions_by_session,
&swarm_plans,
&event_history,
&event_counter,
&swarm_event_tx,
&client_event_tx,
crate::hooks::with_client_terminal_env(
active_terminal_env.clone(),
handle_clear_session(
id,
client_selfdev,
&mut client_session_id,
&client_connection_id,
&agent,
&provider,
&registry,
&sessions,
&shutdown_signals,
&soft_interrupt_queues,
&client_connections,
&swarm_members,
&swarms_by_id,
&file_touch,
&channel_subscriptions,
&channel_subscriptions_by_session,
&swarm_plans,
&event_history,
&event_counter,
&swarm_event_tx,
&client_event_tx,
),
)
.await;
session_control = refresh_session_control_handle(
Expand Down Expand Up @@ -1404,57 +1420,60 @@ pub(super) async fn handle_client(
});
continue;
}
// An empty snapshot is authoritative too. It must clear all
// terminal vars inherited by the daemon rather than retaining
// a prior pane's values.
active_terminal_env = terminal_env;
current_client_instance_id = client_instance_id.clone();
{
let mut connections = client_connections.write().await;
if let Some(info) = connections.get_mut(&client_connection_id) {
info.client_instance_id = client_instance_id.clone();
// Record the client's terminal env so spawn/focus hooks
// target the client's terminal, not the server's stale
// startup env (#405). Only overwrite when the client sent
// something, so reconnects without env don't clobber it.
if !terminal_env.is_empty() {
info.terminal_env = terminal_env.clone();
}
// Record the authoritative snapshot even when it is
// empty so stale pane variables are not retained.
info.terminal_env = active_terminal_env.clone();
}
}
if let Some(target_session_id) = target_session_id {
if crate::session::session_exists(&target_session_id) {
let pre_resume_session_id = client_session_id.clone();
agent = handle_resume_session(
id,
target_session_id.clone(),
subscribe_working_dir.as_deref(),
client_instance_id.as_deref(),
client_has_local_history,
allow_session_takeover,
&mut client_selfdev,
&mut client_session_id,
&client_connection_id,
&agent,
&provider,
&registry,
&sessions,
&shutdown_signals,
&soft_interrupt_queues,
&client_connections,
&client_debug_state,
&swarm_members,
&swarms_by_id,
&file_touch,
&channel_subscriptions,
&channel_subscriptions_by_session,
&swarm_plans,
&swarm_coordinators,
&client_count,
&writer,
&server_name,
&server_icon,
&client_event_tx,
&mcp_pool,
&event_history,
&event_counter,
&swarm_event_tx,
agent = crate::hooks::with_client_terminal_env(
active_terminal_env.clone(),
handle_resume_session(
id,
target_session_id.clone(),
subscribe_working_dir.as_deref(),
client_instance_id.as_deref(),
client_has_local_history,
allow_session_takeover,
&mut client_selfdev,
&mut client_session_id,
&client_connection_id,
&agent,
&provider,
&registry,
&sessions,
&shutdown_signals,
&soft_interrupt_queues,
&client_connections,
&client_debug_state,
&swarm_members,
&swarms_by_id,
&file_touch,
&channel_subscriptions,
&channel_subscriptions_by_session,
&swarm_plans,
&swarm_coordinators,
&client_count,
&writer,
&server_name,
&server_icon,
&client_event_tx,
&mcp_pool,
&event_history,
&event_counter,
&swarm_event_tx,
),
)
.await?;
session_control = refresh_session_control_handle(
Expand Down Expand Up @@ -2753,28 +2772,31 @@ pub(super) async fn handle_client(
}
}

cleanup_client_connection(
&sessions,
&client_session_id,
client_is_processing,
&mut processing_task,
event_handle,
&swarm_members,
&swarms_by_id,
&swarm_coordinators,
&swarm_plans,
&file_touch,
&channel_subscriptions,
&channel_subscriptions_by_session,
&client_debug_state,
&client_debug_id,
&client_connections,
&client_connection_id,
&shutdown_signals,
&soft_interrupt_queues,
&event_history,
&event_counter,
&swarm_event_tx,
crate::hooks::with_client_terminal_env(
active_terminal_env,
cleanup_client_connection(
&sessions,
&client_session_id,
client_is_processing,
&mut processing_task,
event_handle,
&swarm_members,
&swarms_by_id,
&swarm_coordinators,
&swarm_plans,
&file_touch,
&channel_subscriptions,
&channel_subscriptions_by_session,
&client_debug_state,
&client_debug_id,
&client_connections,
&client_connection_id,
&shutdown_signals,
&soft_interrupt_queues,
&event_history,
&event_counter,
&swarm_event_tx,
),
)
.await?;
Ok(())
Expand Down Expand Up @@ -2818,6 +2840,7 @@ async fn start_processing_message(
agent: &Arc<Mutex<Agent>>,
client_event_tx: &mpsc::UnboundedSender<ServerEvent>,
processing_done_tx: &mpsc::UnboundedSender<(u64, Result<()>, Option<String>)>,
client_terminal_env: Vec<(String, String)>,
swarm: &SwarmStatusRefs<'_>,
) {
let ProcessingMessage {
Expand Down Expand Up @@ -2888,32 +2911,34 @@ async fn start_processing_message(
crate::logging::info(&format!("Processing message id={} spawning task", id));
*state.task = Some(tokio::spawn(async move {
let event_tx = tx.clone();
let result = match std::panic::AssertUnwindSafe(process_message_streaming_mpsc(
agent,
&content,
images,
system_reminder,
event_tx,
))
.catch_unwind()
.await
{
Ok(result) => result,
Err(panic_payload) => {
let msg = if let Some(text) = panic_payload.downcast_ref::<&str>() {
text.to_string()
} else if let Some(text) = panic_payload.downcast_ref::<String>() {
text.clone()
} else {
"unknown panic".to_string()
};
crate::logging::error(&format!(
"Processing task PANICKED for message id={}: {}",
id, msg
));
Err(anyhow::anyhow!("Processing task panicked: {}", msg))
}
};
let result =
match std::panic::AssertUnwindSafe(process_message_streaming_mpsc_with_terminal_env(
agent,
&content,
images,
system_reminder,
client_terminal_env,
event_tx,
))
.catch_unwind()
.await
{
Ok(result) => result,
Err(panic_payload) => {
let msg = if let Some(text) = panic_payload.downcast_ref::<&str>() {
text.to_string()
} else if let Some(text) = panic_payload.downcast_ref::<String>() {
text.clone()
} else {
"unknown panic".to_string()
};
crate::logging::error(&format!(
"Processing task PANICKED for message id={}: {}",
id, msg
));
Err(anyhow::anyhow!("Processing task panicked: {}", msg))
}
};
match &result {
Ok(()) => crate::logging::info(&format!(
"Processing task completed OK for message id={}",
Expand Down Expand Up @@ -3255,6 +3280,24 @@ pub(super) async fn process_message_streaming_mpsc(
result
}

/// Process a client-originated turn with that client's terminal identity. The
/// task-local scope flows through turn, pre-tool, and post-tool hooks without
/// mutating process-global environment shared by other client tasks.
async fn process_message_streaming_mpsc_with_terminal_env(
agent: Arc<Mutex<Agent>>,
content: &str,
images: Vec<(String, String)>,
system_reminder: Option<String>,
client_terminal_env: Vec<(String, String)>,
event_tx: tokio::sync::mpsc::UnboundedSender<ServerEvent>,
) -> Result<()> {
crate::hooks::with_client_terminal_env(
client_terminal_env,
process_message_streaming_mpsc(agent, content, images, system_reminder, event_tx),
)
.await
}

#[cfg(test)]
#[path = "client_lifecycle_tests.rs"]
mod tests;
4 changes: 4 additions & 0 deletions crates/jcode-app-core/src/server/client_lifecycle_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -906,6 +906,7 @@ fn reload_starting_rejects_new_turn_without_spawning_processing_task() {
&agent,
&client_event_tx,
&processing_done_tx,
Vec::new(),
&SwarmStatusRefs {
members: &swarm_members,
swarms_by_id: &swarms_by_id,
Expand Down Expand Up @@ -1005,6 +1006,7 @@ async fn client_initiated_turn_fans_out_stream_and_terminal_events_to_live_attac
&agent,
&origin_tx,
&processing_done_tx,
Vec::new(),
&SwarmStatusRefs {
members: &swarm_members,
swarms_by_id: &swarms_by_id,
Expand Down Expand Up @@ -1128,6 +1130,7 @@ fn accepted_reload_recovery_continuation_marks_intent_delivered() -> anyhow::Res
&agent,
&client_event_tx,
&processing_done_tx,
Vec::new(),
&SwarmStatusRefs {
members: &swarm_members,
swarms_by_id: &swarms_by_id,
Expand Down Expand Up @@ -1226,6 +1229,7 @@ fn reload_starting_rejects_new_turns_for_multiple_sessions() {
&agent,
&client_event_tx,
&processing_done_tx,
Vec::new(),
&SwarmStatusRefs {
members: &swarm_members,
swarms_by_id: &swarms_by_id,
Expand Down
Loading
Loading