dpc-tau-ext-shell 0.2.1

//! Runtime state and reader-loop dispatch after the ext-shell handshake.

use crate::tools as path_crate_tools;

#[cfg(test)]
mod tests;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, mpsc};

use tau_proto::{
    CborValue, Event, ExtensionContextReady, HarnessInputMessage, ToolResult, ToolResultKind,
};
use tracing::debug;

use super::{
    DiscoverySourcePolicy, UiShellScheduleContext, UiShellShutdownGenerationCounter,
    apply_started_cwd_metadata, apply_working_directory, cwd_context_event, cwd_notice_event,
    dir_lock_tool_spec, dispatch_action_invoke, dispatch_session_agent_loaded,
    dispatch_session_started, invalid_cwd_context_event, is_shell_tool, schedule_tool_started,
    schedule_ui_shell_command, send_identity_failure, send_ui_shell_saturated_failure,
    with_lock_wait_duration,
};
use crate::Output;
use crate::artifact_transfer::ArtifactTransferManager;
use crate::config::ExtConfig;
use crate::cwd_state::{CwdState, WorkdirSnapshot};
use crate::dir_lock::DirLockManager;
use crate::discovery_source::DiscoverySource;
use crate::scheduler::WorkScheduler;
use crate::tool_lifecycle::{CancelOutcome, ToolCancellationState};

#[derive(Clone, Copy)]
enum StartupCwdSource {
    Process,
    #[cfg(any(test, feature = "echo-agent"))]
    Fixture,
}

pub(super) struct ShellRuntime {
    config: ExtConfig,
    /// Whether session discovery may read process HOME and working-directory
    /// inputs.
    discovery_policy: DiscoverySourcePolicy,
    /// Per-load user snapshots retained while only project roots move.
    discovery_sources: HashMap<tau_proto::AgentId, DiscoverySource>,
    scheduler: Option<WorkScheduler>,
    tx: Output,
    /// Shared pre-effect and active cancellation state for model tool calls.
    cancellation: ToolCancellationState,
    running_ui_commands: Arc<Mutex<HashMap<tau_proto::ShellCommandId, mpsc::Sender<()>>>>,
    shutdown_generation_counter: Arc<UiShellShutdownGenerationCounter>,
    lock_manager: DirLockManager,
    cwd_state: CwdState,
    /// Identifies whether startup cwd comes from the process or a test fixture,
    /// so production freezes the process cwd only after initial configuration.
    startup_cwd_source: StartupCwdSource,
    start_agent_owners: HashMap<String, tau_proto::AgentId>,
    /// Immutable session identity observed from the first lifecycle fact.
    bound_session_id: Option<tau_proto::SessionId>,
    runtime_started: bool,
    /// Main-loop-owned shared artifact transfers.
    artifact_transfers: ArtifactTransferManager,
}

impl ShellRuntime {
    /// Connects worker-side checked-output failures to the manual policy loop.
    pub(super) fn install_waker(&self, waker: tau_client::ManualRuntimeWaker) {
        self.tx.install_waker(waker.clone());
        self.artifact_transfers.install_waker(waker);
    }

    /// Propagates the first checked-output failure observed by a worker.
    pub(super) fn take_mandatory_output_failure(&self) -> tau_client::ClientResult<()> {
        self.tx.take_mandatory_failure()
    }

    /// Creates a production runtime that freezes the process startup cwd during
    /// initial configuration.
    #[cfg(any(test, feature = "echo-agent"))]
    #[cfg_attr(not(test), allow(dead_code))]
    pub(super) fn new(
        tx: Output,
        config: ExtConfig,
        discovery_policy: DiscoverySourcePolicy,
    ) -> Self {
        Self::new_with_artifacts(
            tx,
            ArtifactTransferManager::unavailable(),
            config,
            discovery_policy,
        )
    }

    /// Creates a production runtime with a live Artifact RPC writer.
    pub(super) fn new_with_artifacts(
        tx: Output,
        artifact_transfers: ArtifactTransferManager,
        config: ExtConfig,
        discovery_policy: DiscoverySourcePolicy,
    ) -> Self {
        Self::new_with_cwd_state(
            tx,
            config,
            discovery_policy,
            CwdState::new(),
            StartupCwdSource::Process,
            artifact_transfers,
        )
    }

    #[cfg(any(test, feature = "echo-agent"))]
    /// Creates a fixture runtime with a test-owned startup cwd instead of
    /// reading the process cwd.
    pub(super) fn new_for_test_harness(
        tx: Output,
        config: ExtConfig,
        discovery_policy: DiscoverySourcePolicy,
        fixture_cwd: PathBuf,
    ) -> Self {
        Self::new_with_cwd_state(
            tx,
            config,
            discovery_policy,
            CwdState::new_with_startup_cwd(fixture_cwd),
            StartupCwdSource::Fixture,
            ArtifactTransferManager::unavailable(),
        )
    }

    fn new_with_cwd_state(
        tx: Output,
        config: ExtConfig,
        discovery_policy: DiscoverySourcePolicy,
        cwd_state: CwdState,
        startup_cwd_source: StartupCwdSource,
        artifact_transfers: ArtifactTransferManager,
    ) -> Self {
        Self {
            config,
            discovery_policy,
            discovery_sources: HashMap::new(),
            scheduler: Some(WorkScheduler::new(Default::default())),
            tx,
            cancellation: ToolCancellationState::default(),
            running_ui_commands: Arc::new(Mutex::new(HashMap::new())),
            shutdown_generation_counter: Arc::new(UiShellShutdownGenerationCounter::default()),
            lock_manager: DirLockManager::default(),
            cwd_state,
            startup_cwd_source,
            start_agent_owners: HashMap::new(),
            bound_session_id: None,
            runtime_started: false,
            artifact_transfers,
        }
    }

    fn scheduler(&self) -> tau_client::ClientResult<&WorkScheduler> {
        self.scheduler
            .as_ref()
            .ok_or_else(|| tau_client::ClientError::handler("shell scheduler is shut down"))
    }

    fn send(&self, message: HarnessInputMessage) -> tau_client::ClientResult<()> {
        self.tx.send(message)
    }

    pub(super) fn shutdown(&mut self) {
        self.artifact_transfers.shutdown();
        self.shutdown_generation_counter.advance();
        // Dir-lock waiters must be woken before the scheduler is dropped,
        // because scheduler drop joins workers that may be blocked on locks.
        self.lock_manager.shutdown();
        self.cancellation.lifecycles.prepare_shutdown();
        if let Some(scheduler) = &self.scheduler {
            scheduler.cancel_all_queued();
        }
        let running = self
            .cancellation
            .running_calls
            .lock()
            .expect("running call registry lock poisoned")
            .values()
            .cloned()
            .collect::<Vec<_>>();
        for cancel_tx in running {
            let _ = cancel_tx.send(());
        }
        let running_ui = self
            .running_ui_commands
            .lock()
            .expect("running ui shell registry lock poisoned")
            .values()
            .cloned()
            .collect::<Vec<_>>();
        for cancel_tx in running_ui {
            let _ = cancel_tx.send(());
        }
        self.cwd_state.take_all_pending_workdirs();
    }

    pub(super) fn final_shutdown(&mut self) {
        self.shutdown();
        drop(self.scheduler.take());
        crate::shell_output_spool::shutdown();
    }

    pub(super) fn apply_config(
        &mut self,
        instance_name: tau_proto::ExtensionName,
        tool_prefix: Option<tau_proto::ToolNamePrefix>,
        mut cfg: ExtConfig,
    ) -> tau_client::ClientResult<()> {
        let previous_allowlist_prompt = self.config.shell.allowlist_prompt_fragment();
        if cfg.working_directory.is_none() {
            cfg.working_directory = self.config.working_directory.clone();
        }
        self.cwd_state.set_instance_name(instance_name);
        self.cwd_state.set_context_label(tool_prefix.as_ref());
        if let Err(message) = apply_working_directory(&self.config, &cfg, self.runtime_started) {
            return Err(tau_client::ClientError::handler(message));
        }
        if matches!(self.startup_cwd_source, StartupCwdSource::Process) {
            self.cwd_state
                .freeze_process_startup_cwd()
                .map_err(tau_client::ClientError::handler)?;
        }
        if let Err(message) = self.lock_manager.configure(&cfg.dir_lock) {
            return Err(tau_client::ClientError::handler(message));
        }

        let dir_lock_was_enabled = self.config.dir_lock.enable;
        let dir_lock_changed = dir_lock_was_enabled != cfg.dir_lock.enable;
        let dir_lock_disabling = dir_lock_was_enabled && !cfg.dir_lock.enable;
        self.config = cfg;
        if dir_lock_disabling {
            let _ = self.lock_manager.disable();
        }
        if dir_lock_changed {
            self.tx
                .register_local_tool(tau_proto::ToolRegistrationDeclared {
                    tool: dir_lock_tool_spec(self.config.dir_lock.enable),
                    tool_group: Some(tau_proto::ToolGroup {
                        name: tau_proto::ToolGroupName::new("shell"),
                        prompt_fragment: None,
                    }),
                    prompt_fragment: None,
                })?;
        }
        if previous_allowlist_prompt != self.config.shell.allowlist_prompt_fragment() {
            self.send(HarnessInputMessage::emit_with_persist(
                Event::ExtPromptFragmentPublish(tau_proto::ExtPromptFragmentPublish {
                    fragment: crate::shell_workdir_prompt_fragment(&self.config.shell),
                }),
                false,
            ))?;
        }
        tracing::info!(
            target: "tau_ext_shell",
            directory_locking = self.config.dir_lock.enable,
            "shell configured"
        );
        Ok(())
    }

    pub(super) fn handle_event(
        &mut self,
        event: Event,
        is_replay: bool,
    ) -> tau_client::ClientResult<()> {
        self.runtime_started = true;
        match event {
            Event::AgentStarted(started) => {
                apply_started_cwd_metadata(started, &self.tx, &self.cwd_state, is_replay)?;
            }
            Event::ToolStarted(invoke) => {
                let local_tool_name = invoke.tool_name.clone();
                self.handle_tool_started(invoke, &local_tool_name, is_replay)?;
            }
            Event::SessionStarted(started) => {
                match self.bound_session_id.as_ref() {
                    Some(bound) if bound != &started.session_id => {
                        return Err(tau_client::ClientError::handler(format!(
                            "immutable session mismatch: expected `{bound}`, received `{}`",
                            started.session_id
                        )));
                    }
                    Some(_) => return Ok(()),
                    None => self.bound_session_id = Some(started.session_id.clone()),
                }
                self.artifact_transfers
                    .bind_session(started.session_id.clone());
                dispatch_session_started(started, &self.tx, self.discovery_policy)?;
            }
            Event::SessionAgentLoaded(loaded) => {
                dispatch_session_agent_loaded(
                    loaded,
                    &self.tx,
                    &self.cwd_state,
                    true,
                    self.discovery_policy,
                )?;
            }
            Event::SessionAgentUnloaded(unloaded) => {
                if !is_replay {
                    self.handle_session_agent_unloaded(unloaded);
                }
            }
            Event::AgentMetadataSet(set) => self.handle_agent_metadata_set(set, is_replay)?,
            Event::AgentMetadataUnset(unset) => {
                self.handle_agent_metadata_unset(unset, is_replay)?
            }
            Event::AgentReplayComplete(done) => self.handle_agent_replay_complete(done)?,
            Event::HarnessAgentDiscoveryRefreshRequested(request) if !is_replay => {
                self.handle_discovery_refresh(request)?;
            }
            Event::HarnessAgentContextInitialized(context) if !is_replay => {
                self.handle_discovery_installed(context)?;
            }
            Event::SessionShutdown(shutdown) => {
                if self.bound_session_id.as_ref() != Some(&shutdown.session_id) {
                    return Err(tau_client::ClientError::handler(
                        "session shutdown does not match immutable binding",
                    ));
                }
                self.shutdown_session();
            }
            Event::StartAgentAccepted(accepted) => {
                self.start_agent_owners
                    .insert(accepted.query_id, accepted.agent_id);
            }
            Event::StartAgentResult(result) => self.handle_start_agent_result(result),
            Event::ActionInvoke(invoke) => {
                #[cfg(test)]
                if invoke.action_id == "test.saturate-optional-output" {
                    for _ in 0..70 {
                        let _ = self.send(HarnessInputMessage::emit(dispatch_action_invoke(
                            invoke.clone(),
                            &self.lock_manager,
                        )));
                    }
                    return Ok(());
                }
                self.send(HarnessInputMessage::emit(dispatch_action_invoke(
                    invoke,
                    &self.lock_manager,
                )))?;
            }
            Event::ToolCancelRequest(request) => self.handle_tool_cancel_request(request),
            Event::UiShellCommand(cmd) => self.handle_ui_shell_command(cmd)?,
            _ => {}
        }
        Ok(())
    }

    fn handle_tool_started(
        &self,
        invoke: tau_proto::ToolStarted,
        local_tool_name: &tau_proto::ToolName,
        is_replay: bool,
    ) -> tau_client::ClientResult<()> {
        if is_replay || !is_shell_tool(local_tool_name.as_str()) {
            return Ok(());
        }
        if let Err(error) = schedule_tool_started(
            (invoke, local_tool_name),
            self.scheduler()?,
            &self.tx,
            self.config.clone(),
            self.lock_manager.clone(),
            self.cancellation.clone(),
            self.cwd_state.clone(),
            self.artifact_transfers.control(),
        ) {
            let (identity, failure) = *error;
            let _ = send_identity_failure(identity, failure, &self.tx);
        }
        Ok(())
    }

    pub(super) fn handle_scoped_tool_started(
        &mut self,
        invoke: tau_proto::ToolStarted,
        local_tool_name: &tau_proto::ToolName,
    ) -> tau_client::ClientResult<()> {
        self.runtime_started = true;
        self.handle_tool_started(invoke, local_tool_name, false)
    }

    fn handle_session_agent_unloaded(&mut self, unloaded: tau_proto::SessionAgentUnloaded) {
        self.lock_manager.release_agent(&unloaded.agent_id);
        if let Some(scheduler) = &self.scheduler {
            scheduler.cancel_agent(&unloaded.agent_id);
        }
        self.cancellation
            .lifecycles
            .remove_agent(&unloaded.agent_id);
        self.cwd_state.unset(&unloaded.agent_id);
        self.cwd_state.take_pending_ready(&unloaded.agent_id);
        self.cwd_state.remove_initialization(&unloaded.agent_id);
        self.discovery_sources.remove(&unloaded.agent_id);
        self.cwd_state
            .take_pending_workdir_result(&unloaded.agent_id);
        self.start_agent_owners
            .retain(|_, agent_id| agent_id != &unloaded.agent_id);
    }

    fn handle_agent_metadata_set(
        &mut self,
        set: tau_proto::AgentMetadataSet,
        is_replay: bool,
    ) -> tau_client::ClientResult<()> {
        if set.key != self.cwd_state.key() {
            return Ok(());
        }
        if self.cwd_state.is_replay_failed(&set.agent_id) {
            return Ok(());
        }
        if let CborValue::Text(path) = set.value {
            self.handle_text_cwd_metadata_set(
                set.agent_id,
                PathBuf::from(path),
                set.mutation_id.as_ref(),
                is_replay,
            )
        } else {
            self.handle_invalid_cwd_metadata_set(set.agent_id, set.mutation_id.as_ref(), is_replay)
        }
    }

    fn handle_text_cwd_metadata_set(
        &mut self,
        agent_id: tau_proto::AgentId,
        cwd: PathBuf,
        mutation_id: Option<&tau_proto::AgentMetadataMutationId>,
        is_replay: bool,
    ) -> tau_client::ClientResult<()> {
        if !self
            .cwd_state
            .set_metadata_text(agent_id.clone(), cwd.clone())
        {
            return self.handle_invalid_cwd_metadata_set(agent_id, mutation_id, is_replay);
        }
        if is_replay {
            return Ok(());
        }
        if let Some((session_id, initialization_id)) = self.cwd_state.initialization(&agent_id) {
            self.tx
                .send_checked(HarnessInputMessage::emit_transient(cwd_context_event(
                    session_id,
                    agent_id.clone(),
                    initialization_id,
                    &cwd,
                    &self.cwd_state,
                )))?;
        }
        let pending_workdir =
            self.cwd_state
                .committed_pending_workdir_result(&agent_id, &cwd, mutation_id);
        if pending_workdir.is_some() {
            let _ = self.tx.send(HarnessInputMessage::emit(cwd_notice_event(
                agent_id.clone(),
                &cwd,
            )));
        }
        // Loaded agents settle only after the correlated replacement is
        // durable. The old direct path remains for standalone
        // fixture/tool lifecycles.
        if !self.discovery_sources.contains_key(&agent_id) {
            self.complete_pending_workdir_after_text_metadata(pending_workdir, &cwd)?;
        }
        self.publish_ready_if_pending(agent_id)
    }

    fn complete_pending_workdir_after_text_metadata(
        &self,
        pending_workdir: Option<crate::cwd_state::CompletedPendingWorkdir>,
        cwd: &Path,
    ) -> tau_client::ClientResult<()> {
        if let Some(pending_workdir) = pending_workdir {
            let call_id = pending_workdir.identity.call_id.clone();
            let event = if pending_workdir.cancel_requested {
                Event::ToolCancelled(tau_proto::ToolCancelled {
                    presentation: Default::default(),
                    call_id: pending_workdir.identity.call_id,
                    tool_name: pending_workdir.identity.wire_tool_name,
                    tool_type: tau_proto::ToolType::Function,
                    display: None,
                })
            } else if pending_workdir.matched_request {
                let output = path_crate_tools::workdir::output(cwd);
                Event::ToolResult(ToolResult {
                    presentation: Default::default(),
                    call_id: pending_workdir.identity.call_id,
                    tool_name: pending_workdir.identity.wire_tool_name,
                    tool_type: tau_proto::ToolType::Function,
                    result: output.result,
                    provider_content: Vec::new(),
                    kind: ToolResultKind::Final,
                    display: Some(output.display),
                    originator: pending_workdir.identity.originator,
                })
            } else {
                Event::ToolError(tau_proto::ToolError {
                    presentation: Default::default(),
                    call_id: pending_workdir.identity.call_id,
                    tool_name: pending_workdir.identity.wire_tool_name,
                    tool_type: tau_proto::ToolType::Function,
                    message: format!(
                        "committed cwd metadata did not match requested cwd; cwd changed to {}",
                        cwd.display()
                    ),
                    details: None,
                    display: None,
                    originator: pending_workdir.identity.originator,
                })
            };
            self.tx.report_tool_terminal(with_lock_wait_duration(
                event,
                pending_workdir.lock_wait_duration_seconds,
            ))?;
            self.cwd_state.take_pending_workdir_by_call(&call_id);
        }
        Ok(())
    }

    fn handle_invalid_cwd_metadata_set(
        &mut self,
        agent_id: tau_proto::AgentId,
        mutation_id: Option<&tau_proto::AgentMetadataMutationId>,
        is_replay: bool,
    ) -> tau_client::ClientResult<()> {
        self.cwd_state.set_invalid(agent_id.clone());
        if is_replay {
            return Ok(());
        }
        if let Some((session_id, initialization_id)) = self.cwd_state.initialization(&agent_id) {
            self.tx.send_checked(HarnessInputMessage::emit_transient(
                invalid_cwd_context_event(
                    session_id,
                    agent_id.clone(),
                    initialization_id,
                    &self.cwd_state,
                ),
            ))?;
        }
        if !self.discovery_sources.contains_key(&agent_id)
            && let Some(pending) = self
                .cwd_state
                .correlated_pending_workdir_result(&agent_id, mutation_id)
        {
            self.send_pending_workdir_error(
                pending,
                "committed workdir metadata is malformed; workdir setter was superseded",
            )?;
        }
        self.publish_ready_if_pending(agent_id)
    }

    fn handle_agent_metadata_unset(
        &mut self,
        unset: tau_proto::AgentMetadataUnset,
        is_replay: bool,
    ) -> tau_client::ClientResult<()> {
        if unset.key != self.cwd_state.key() {
            return Ok(());
        }
        if self.cwd_state.is_replay_failed(&unset.agent_id) {
            return Ok(());
        }
        self.cwd_state.unset(&unset.agent_id);
        if is_replay {
            return Ok(());
        }
        if let Ok(cwd) = self.cwd_state.process_default() {
            self.tx.send_checked(HarnessInputMessage::emit_transient(
                Event::AgentMetadataSetRequest(tau_proto::AgentMetadataSet {
                    agent_id: unset.agent_id,
                    key: self.cwd_state.key(),
                    value: CborValue::Text(cwd.display().to_string()),
                    mutation_id: None,
                    inheritable: true,
                }),
            ))?;
        }
        Ok(())
    }

    fn send_pending_workdir_error(
        &self,
        pending: crate::cwd_state::CompletedPendingWorkdir,
        message: &str,
    ) -> tau_client::ClientResult<()> {
        let call_id = pending.identity.call_id.clone();
        let event = if pending.cancel_requested {
            Event::ToolCancelled(tau_proto::ToolCancelled {
                presentation: Default::default(),
                call_id: pending.identity.call_id,
                tool_name: pending.identity.wire_tool_name,
                tool_type: tau_proto::ToolType::Function,
                display: None,
            })
        } else {
            Event::ToolError(tau_proto::ToolError {
                presentation: Default::default(),
                call_id: pending.identity.call_id,
                tool_name: pending.identity.wire_tool_name,
                tool_type: tau_proto::ToolType::Function,
                message: message.to_owned(),
                details: None,
                display: None,
                originator: pending.identity.originator,
            })
        };
        self.tx.report_tool_terminal(event)?;
        self.cwd_state.take_pending_workdir_by_call(&call_id);
        Ok(())
    }

    fn publish_ready_if_pending(
        &self,
        agent_id: tau_proto::AgentId,
    ) -> tau_client::ClientResult<()> {
        if let Some((session_id, agent_initialization_id)) =
            self.cwd_state.take_pending_ready(&agent_id)
        {
            self.tx.send_checked(HarnessInputMessage::emit_transient(
                Event::ExtensionContextReady(ExtensionContextReady {
                    session_id,
                    agent_id,
                    agent_initialization_id,
                }),
            ))?;
        }
        Ok(())
    }

    fn handle_agent_replay_complete(
        &mut self,
        done: tau_proto::AgentReplayComplete,
    ) -> tau_client::ClientResult<()> {
        let Some((session_id, initialization_id)) = self.cwd_state.pending_ready(&done.agent_id)
        else {
            return Ok(());
        };
        if done.error.is_some() {
            self.cwd_state.take_pending_ready(&done.agent_id);
            self.cwd_state.set_replay_failed(done.agent_id);
            return Ok(());
        }
        let source = if self.discovery_policy.reads_environment() {
            DiscoverySource::new(dirs::home_dir())
        } else {
            DiscoverySource::empty()
        };
        self.discovery_sources.insert(done.agent_id.clone(), source);
        let cwd = match self.cwd_state.snapshot(&done.agent_id) {
            Ok(WorkdirSnapshot::Valid(cwd)) => Ok(cwd),
            Ok(_) => Err("remembered cwd is malformed or replay failed".to_owned()),
            Err(error) => Err(error),
        };
        self.publish_current_discovery(
            session_id.clone(),
            done.agent_id.clone(),
            initialization_id.clone(),
            cwd,
            None,
        )?;
        if let Some(cwd) = self.cwd_state.get(&done.agent_id) {
            self.tx
                .send_checked(HarnessInputMessage::emit_transient(cwd_context_event(
                    session_id.clone(),
                    done.agent_id.clone(),
                    initialization_id.clone(),
                    &cwd,
                    &self.cwd_state,
                )))?;
            return self.publish_ready_if_pending(done.agent_id);
        }
        if self.cwd_state.is_invalid(&done.agent_id) {
            self.tx.send_checked(HarnessInputMessage::emit_transient(
                invalid_cwd_context_event(
                    session_id.clone(),
                    done.agent_id.clone(),
                    initialization_id.clone(),
                    &self.cwd_state,
                ),
            ))?;
            return self.publish_ready_if_pending(done.agent_id);
        }
        let Ok(cwd) = self.cwd_state.process_default() else {
            self.cwd_state.take_pending_ready(&done.agent_id);
            return Ok(());
        };
        self.tx.send_checked(HarnessInputMessage::emit_transient(
            Event::AgentMetadataSetRequest(tau_proto::AgentMetadataSet {
                agent_id: done.agent_id.clone(),
                key: self.cwd_state.key(),
                value: CborValue::Text(cwd.display().to_string()),
                mutation_id: None,
                inheritable: true,
            }),
        ))?;
        self.cwd_state
            .set_pending_ready(done.agent_id, session_id, initialization_id);
        Ok(())
    }

    /// Scan on the shell execution host, retaining only user inputs on failure.
    fn publish_current_discovery(
        &self,
        session_id: tau_proto::SessionId,
        agent_id: tau_proto::AgentId,
        initialization_id: tau_proto::AgentInitializationId,
        cwd: Result<PathBuf, String>,
        refresh_id: Option<u64>,
    ) -> tau_client::ClientResult<()> {
        let Some(source) = self.discovery_sources.get(&agent_id) else {
            return Ok(());
        };
        let cwd = cwd.and_then(|cwd| {
            if !cwd.is_absolute() {
                return Err("remembered cwd is not absolute".to_owned());
            }
            std::fs::read_dir(&cwd).map_err(|error| {
                format!("project discovery at {} failed: {error}", cwd.display())
            })?;
            Ok(cwd)
        });
        let discovery_error = cwd.as_ref().err().cloned();
        let project = if self.discovery_policy.reads_environment() {
            cwd.as_ref().ok().map(PathBuf::as_path)
        } else {
            None
        };
        let scan = source.scan_project(session_id.clone(), project);
        let binding = refresh_id.is_none().then(|| {
            let user = source.scan_project(session_id.clone(), None).snapshot;
            tau_proto::DiscoveryWorkdirBinding {
                metadata_key: self.cwd_state.key(),
                user_skills: user.skills,
                user_candidates: source.user_candidates(),
                retained_user_state: source.retained_user_state(),
                user_agents_files: user.agents_files,
            }
        });
        self.tx.send_checked(HarnessInputMessage::emit_transient(
            Event::ExtensionAgentDiscoverySnapshotDeclared(
                tau_proto::ExtensionAgentDiscoverySnapshotDeclared {
                    workdir_binding: binding,
                    refresh_id,
                    discovery_error,
                    frontmatter_diagnostics: scan.snapshot.frontmatter_diagnostics,
                    session_id,
                    agent_id,
                    agent_initialization_id: initialization_id,
                    skills: scan.snapshot.skills,
                    agents_files: scan.snapshot.agents_files,
                },
            ),
        ))
    }

    /// Reply only to this loaded source's canonical committed-value request.
    fn handle_discovery_refresh(
        &mut self,
        request: tau_proto::HarnessAgentDiscoveryRefreshRequested,
    ) -> tau_client::ClientResult<()> {
        if request.metadata_key != self.cwd_state.key()
            || self.cwd_state.initialization(&request.agent_id)
                != Some((
                    request.session_id.clone(),
                    request.agent_initialization_id.clone(),
                ))
        {
            return Ok(());
        }
        let restored = DiscoverySource::from_retained_user_state(&request.retained_user_state);
        let cwd = match request.metadata_value {
            Some(CborValue::Text(path)) => Ok(PathBuf::from(path)),
            _ => Err("committed cwd is absent or malformed; retained user context only".to_owned()),
        }
        .and_then(|cwd| restored.as_ref().map(|_| cwd).map_err(Clone::clone));
        if let Ok(source) = restored {
            self.discovery_sources
                .insert(request.agent_id.clone(), source);
        }
        self.publish_current_discovery(
            request.session_id,
            request.agent_id,
            request.agent_initialization_id,
            cwd,
            Some(request.refresh_id),
        )
    }

    /// Release a setter only after its matching replacement was installed.
    fn handle_discovery_installed(
        &self,
        context: tau_proto::HarnessAgentContextInitialized,
    ) -> tau_client::ClientResult<()> {
        if self.cwd_state.initialization(&context.agent_id)
            != Some((context.session_id, context.agent_initialization_id))
        {
            return Ok(());
        }
        for refresh in context.discovery_refreshes {
            if refresh.metadata_key != self.cwd_state.key() {
                continue;
            }
            let Some(pending) = self
                .cwd_state
                .correlated_pending_workdir_result(&context.agent_id, refresh.mutation_id.as_ref())
            else {
                continue;
            };
            if let Some(error) = refresh.error {
                self.send_pending_workdir_error(
                    pending,
                    &format!("cwd metadata committed, but discovery failed: {error}"),
                )?;
            } else if let Some(cwd) = self.cwd_state.get(&context.agent_id) {
                let pending = self.cwd_state.committed_pending_workdir_result(
                    &context.agent_id,
                    &cwd,
                    refresh.mutation_id.as_ref(),
                );
                self.complete_pending_workdir_after_text_metadata(pending, &cwd)?;
            } else {
                self.send_pending_workdir_error(
                    pending,
                    "cwd metadata committed, but current cwd is unavailable",
                )?;
            }
        }
        Ok(())
    }

    fn shutdown_session(&mut self) {
        self.shutdown();
        self.start_agent_owners.clear();
        self.discovery_sources.clear();
    }

    fn handle_start_agent_result(&mut self, result: tau_proto::StartAgentResult) {
        if let Some(agent_id) = self.start_agent_owners.remove(&result.query_id) {
            self.lock_manager.release_agent(&agent_id);
            if let Some(scheduler) = &self.scheduler {
                scheduler.cancel_agent(&agent_id);
            }
            self.cancellation.lifecycles.remove_agent(&agent_id);
        }
    }

    fn handle_tool_cancel_request(&mut self, request: tau_proto::ToolCancelRequest) {
        if self
            .cwd_state
            .request_pending_workdir_cancel(&request.target_call_id)
        {
            debug!(call_id = %request.target_call_id, "workdir cancellation deferred to metadata commit");
            return;
        }
        let Some(outcome) = self.cancellation.lifecycles.cancel(&request.target_call_id) else {
            debug!(call_id = %request.target_call_id, "tool cancellation requested for unknown call");
            return;
        };
        if outcome == CancelOutcome::PreventedEffect {
            if let Some(scheduler) = &self.scheduler {
                scheduler.cancel_queued_call(&request.target_call_id);
            }
            self.lock_manager
                .cancel_waiting_call(&request.target_call_id);
            debug!(call_id = %request.target_call_id, "tool cancellation prevented effect start");
            return;
        }
        self.artifact_transfers
            .cancel(&request.target_call_id, &self.tx);
        let cancel_tx = self
            .cancellation
            .running_calls
            .lock()
            .expect("running call registry lock poisoned")
            .get(&request.target_call_id)
            .cloned();
        if let Some(cancel_tx) = cancel_tx {
            debug!(call_id = %request.target_call_id, "tool cancellation requested for running call");
            if cancel_tx.send(()).is_err() {
                debug!(call_id = %request.target_call_id, "shell cancellation receiver already gone");
            }
        } else {
            if self
                .lock_manager
                .cancel_waiting_call(&request.target_call_id)
            {
                debug!(call_id = %request.target_call_id, "cancellation requested for active dir-lock waiter");
            } else {
                debug!(call_id = %request.target_call_id, "active cancellation recorded before sender registration");
            }
        }
    }

    /// Drains artifact filesystem completions before the manual loop sleeps.
    pub(super) fn drain_artifact_commands(&mut self) -> tau_client::ClientResult<()> {
        let scheduler = self
            .scheduler
            .as_ref()
            .ok_or_else(|| tau_client::ClientError::handler("shell scheduler is shut down"))?;
        self.artifact_transfers.drain(scheduler, &self.tx);
        Ok(())
    }

    /// Routes one exact-correlated Artifact response to its typed transfer.
    pub(super) fn handle_artifact_result(
        &mut self,
        result: tau_proto::ArtifactResult,
    ) -> tau_client::ClientResult<()> {
        let scheduler = self
            .scheduler
            .as_ref()
            .ok_or_else(|| tau_client::ClientError::handler("shell scheduler is shut down"))?;
        self.artifact_transfers
            .handle_result(result, scheduler, &self.tx);
        Ok(())
    }

    fn handle_ui_shell_command(
        &self,
        cmd: tau_proto::UiShellCommand,
    ) -> tau_client::ClientResult<()> {
        if cmd
            .target_agent_id
            .as_ref()
            .is_some_and(|agent_id| self.cwd_state.pending_ready(agent_id).is_some())
        {
            send_ui_shell_saturated_failure(
                cmd,
                "workdir replay is not complete for the target agent".to_owned(),
                &self.tx,
            );
            return Ok(());
        }
        let cwd = match cmd.target_agent_id.as_ref() {
            Some(agent_id) => self.cwd_state.get_or_default(agent_id),
            None => self.cwd_state.process_default(),
        };
        let cwd = match cwd {
            Ok(cwd) => cwd,
            Err(message) => {
                send_ui_shell_saturated_failure(cmd, message, &self.tx);
                return Ok(());
            }
        };
        if let Err(error) = schedule_ui_shell_command(
            cmd,
            UiShellScheduleContext {
                scheduler: self.scheduler()?,
                tx: &self.tx,
                shell_config: self.config.shell.clone(),
                running_ui_commands: Arc::clone(&self.running_ui_commands),
                shutdown_generation_counter: Arc::clone(&self.shutdown_generation_counter),
                scheduled_generation: self.shutdown_generation_counter.current(),
                cwd,
            },
        ) {
            let (cmd, message) = *error;
            send_ui_shell_saturated_failure(cmd, message, &self.tx);
        }
        Ok(())
    }
}