winx-code-agent 0.2.358

use std::path::PathBuf;

use rmcp::model::CallToolRequestParams;

use super::principal::{
    canonical_workspace_identity, expected_affinity_thread_id, expected_scratch_thread_id,
    RequestScope,
};
use super::{SessionIsolation, WinxService};
use crate::config::HttpSessionAffinity;
use crate::errors::{Result, WinxError};
use crate::tool_registry::ToolKind;

/// Result recorded in usage telemetry after the request's project identity has
/// been checked. It contains no user path or transport identity.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum WorkspaceCoherence {
    NotRequired,
    FirstCall,
    Validated,
}

impl WorkspaceCoherence {
    pub(super) const fn as_str(self) -> &'static str {
        match self {
            Self::NotRequired => "not_required",
            Self::FirstCall => "first_call",
            Self::Validated => "validated",
        }
    }
}

impl WinxService {
    /// Fail closed before a remote tool can select a shell or touch a file.
    ///
    /// `workspace_root` is deliberately not compared with target file paths.
    /// It is the project identity paired with `thread_id`; filesystem reach is
    /// still governed independently by `WINX_ALLOW_PATHS` and mode policy.
    pub(super) async fn validate_workspace_coherence(
        &self,
        request: &CallToolRequestParams,
        scope: &RequestScope,
        affinity: HttpSessionAffinity,
        conversation_identity: Option<&str>,
    ) -> Result<WorkspaceCoherence> {
        if self.isolation != SessionIsolation::Strict {
            return Ok(WorkspaceCoherence::NotRequired);
        }

        let arguments = request.arguments.as_ref();
        let internal_thread_id = string_argument(arguments, "thread_id").unwrap_or_default();
        let external_thread_id = scope.unscope_text(&internal_thread_id);

        if ToolKind::parse(request.name.as_ref()) == Some(ToolKind::Initialize) {
            return self
                .validate_remote_initialize(
                    arguments,
                    &internal_thread_id,
                    &external_thread_id,
                    affinity,
                    conversation_identity,
                )
                .await;
        }

        let workspace_root = string_argument(arguments, "workspace_root").ok_or_else(|| {
            WinxError::WorkspaceBindingRequired { thread_id: external_thread_id.clone() }
        })?;
        if external_thread_id.is_empty() {
            return Err(WinxError::WorkspaceBindingRequired { thread_id: String::new() });
        }
        let requested_workspace = canonical_binding_root(&workspace_root)?;

        // Ids derived from the workspace path must match it. Scratch and
        // saved-task resumes derive their id before the final workspace is
        // known (playground, snapshot), so for them the established binding
        // below is what ties the id to the workspace Initialize returned.
        let affinity_mismatch = affinity_thread_mismatch(
            &requested_workspace,
            &external_thread_id,
            affinity,
            conversation_identity,
        );

        let bound_workspace =
            if let Some(workspace) = self.bound_workspace(&internal_thread_id).await {
                workspace
            } else if affinity_mismatch {
                // No binding vouches for this id, so it was invented or taken
                // from another workspace.
                return Err(WinxError::WorkspaceThreadMismatch {
                    thread_id: external_thread_id,
                    workspace_root: requested_workspace,
                });
            } else {
                // After an adapter restart the registry is empty even though the
                // session's snapshot (and possibly its guardian PTY) survived.
                // Rehydrate before rejecting so a remote tool call carrying a
                // known thread_id/workspace pair keeps working across restarts.
                drop(self.tool_session_for(&internal_thread_id).await);
                self.bound_workspace(&internal_thread_id)
                    .await
                    .ok_or(WinxError::BashStateNotInitialized)?
            };
        let bound_workspace = canonical_workspace_identity(&bound_workspace.to_string_lossy());
        if requested_workspace != bound_workspace {
            if affinity_mismatch {
                // The id belongs to another workspace entirely.
                return Err(WinxError::WorkspaceThreadMismatch {
                    thread_id: external_thread_id,
                    workspace_root: requested_workspace,
                });
            }
            return Err(WinxError::WorkspaceBindingMismatch {
                thread_id: external_thread_id,
                requested_workspace,
                bound_workspace,
            });
        }

        Ok(WorkspaceCoherence::Validated)
    }

    async fn validate_remote_initialize(
        &self,
        arguments: Option<&rmcp::model::JsonObject>,
        internal_thread_id: &str,
        external_thread_id: &str,
        affinity: HttpSessionAffinity,
        conversation_identity: Option<&str>,
    ) -> Result<WorkspaceCoherence> {
        let kind = string_argument(arguments, "type").unwrap_or_else(|| "first_call".to_string());
        let workspace = string_argument(arguments, "any_workspace_path").unwrap_or_default();

        if kind == "user_asked_change_workspace" {
            let workspace_root = if workspace.is_empty() {
                PathBuf::new()
            } else {
                canonical_workspace_identity(&workspace)
            };
            return Err(WinxError::WorkspaceChangeRequiresNewSession { workspace_root });
        }

        // An empty first-call path intentionally creates a per-thread scratch
        // playground. There is no canonical input root to compare before it is
        // created; Initialize returns the concrete binding for later calls.
        if kind == "first_call" && workspace.trim().is_empty() {
            return Ok(WorkspaceCoherence::FirstCall);
        }

        let requested_workspace = canonical_initialize_workspace(&workspace)?;
        let mut bound_workspace = self.bound_workspace(internal_thread_id).await;
        // Stateless clients probe a known pair with `status` after a gap. After
        // an adapter restart the registry is empty although the snapshot
        // survived, so rehydrate exactly like the other tools do instead of
        // forcing a full first_call. An id that does not derive from this
        // workspace never loads a snapshot and stays an unknown session.
        if bound_workspace.is_none()
            && kind == "status"
            && !affinity_thread_mismatch(
                &requested_workspace,
                external_thread_id,
                affinity,
                conversation_identity,
            )
        {
            drop(self.tool_session_for(internal_thread_id).await);
            bound_workspace = self.bound_workspace(internal_thread_id).await;
        }
        let mut recovering_missing_session = false;
        if let Some(bound_workspace) = bound_workspace {
            let bound_workspace = canonical_workspace_identity(&bound_workspace.to_string_lossy());
            if requested_workspace != bound_workspace {
                return Err(WinxError::WorkspaceBindingMismatch {
                    thread_id: external_thread_id.to_string(),
                    requested_workspace,
                    bound_workspace,
                });
            }
        } else if kind != "first_call" {
            if matches!(kind.as_str(), "reset_shell" | "user_asked_mode_change") {
                recovering_missing_session = true;
            } else {
                return Err(WinxError::BashStateNotInitialized);
            }
        }

        Ok(if kind == "first_call" || recovering_missing_session {
            WorkspaceCoherence::FirstCall
        } else {
            WorkspaceCoherence::Validated
        })
    }
}

/// Whether a workspace-derived affinity id disagrees with the supplied one.
/// Scratch ids are derived before their playground exists, so they are
/// vouched for by the established binding instead.
fn affinity_thread_mismatch(
    requested_workspace: &std::path::Path,
    external_thread_id: &str,
    affinity: HttpSessionAffinity,
    conversation_identity: Option<&str>,
) -> bool {
    expected_affinity_thread_id(requested_workspace, affinity, conversation_identity).is_some_and(
        |expected| {
            external_thread_id != expected
                && expected_scratch_thread_id(affinity, conversation_identity).as_deref()
                    != Some(external_thread_id)
        },
    )
}

fn canonical_binding_root(workspace: &str) -> Result<PathBuf> {
    let supplied = absolute_workspace_path(workspace, "workspace_root")?;
    if supplied.is_file() {
        return Err(WinxError::ParameterValidationError {
            field: "workspace_root".to_string(),
            message: "must be the workspace directory returned by Initialize, not a file path"
                .to_string(),
        });
    }
    let path = canonical_workspace_identity(workspace);
    Ok(path)
}

fn canonical_initialize_workspace(workspace: &str) -> Result<PathBuf> {
    let _ = absolute_workspace_path(workspace, "any_workspace_path")?;
    // Initialize intentionally accepts either a directory or a file hint; a
    // file resolves to its parent exactly as the tool implementation does.
    Ok(canonical_workspace_identity(workspace))
}

fn absolute_workspace_path(workspace: &str, field: &str) -> Result<PathBuf> {
    if workspace.trim().is_empty() {
        return Err(WinxError::ParameterValidationError {
            field: field.to_string(),
            message: "must be a non-empty absolute workspace path".to_string(),
        });
    }
    let expanded = crate::utils::path::expand_user(workspace);
    let supplied = PathBuf::from(expanded);
    if !supplied.is_absolute() {
        return Err(WinxError::ParameterValidationError {
            field: field.to_string(),
            message: "must be an absolute path for a remote session".to_string(),
        });
    }
    Ok(supplied)
}

fn string_argument(arguments: Option<&rmcp::model::JsonObject>, key: &str) -> Option<String> {
    arguments?
        .get(key)?
        .as_str()
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToString::to_string)
}

#[cfg(test)]
mod tests {
    #![allow(clippy::expect_used)]

    use super::*;
    use std::path::Path;

    use crate::state::BashState;
    use serde_json::Value;

    fn request(name: &str, value: &Value) -> CallToolRequestParams {
        CallToolRequestParams::new(name.to_string())
            .with_arguments(value.as_object().expect("request object").clone())
    }

    async fn bind(service: &WinxService, thread_id: &str, workspace: &Path) {
        let (slot, guard) = service.session_for(thread_id).await;
        let mut state = BashState::new();
        state.current_thread_id = thread_id.to_string();
        state.workspace_root = workspace.to_path_buf();
        state.cwd = workspace.to_path_buf();
        state.initialized = true;
        *slot.lock().await = Some(state);
        drop(guard);
    }

    #[tokio::test]
    async fn strict_remote_binding_allows_targets_outside_the_workspace() {
        let workspace = tempfile::tempdir().expect("workspace");
        let external = tempfile::tempdir().expect("external target");
        let service = WinxService::with_isolation(SessionIsolation::Strict);
        bind(&service, "thread", workspace.path()).await;
        let request = request(
            "ReadFiles",
            &serde_json::json!({
                "file_paths": [external.path().join("outside.txt")],
                "thread_id": "thread",
                "workspace_root": workspace.path()
            }),
        );

        let result = service
            .validate_workspace_coherence(
                &request,
                &RequestScope::default(),
                HttpSessionAffinity::Thread,
                None,
            )
            .await;
        assert_eq!(result.expect("coherent binding"), WorkspaceCoherence::Validated);
    }

    #[tokio::test]
    async fn strict_remote_binding_rejects_another_projects_thread() {
        let intended = tempfile::tempdir().expect("intended workspace");
        let wrong = tempfile::tempdir().expect("wrong workspace");
        let service = WinxService::with_isolation(SessionIsolation::Strict);
        bind(&service, "wrong_thread", wrong.path()).await;
        let request = request(
            "BashCommand",
            &serde_json::json!({
                "action_json": {"type": "command", "command": "pwd"},
                "thread_id": "wrong_thread",
                "workspace_root": intended.path()
            }),
        );

        let error = service
            .validate_workspace_coherence(
                &request,
                &RequestScope::default(),
                HttpSessionAffinity::Thread,
                None,
            )
            .await
            .expect_err("cross-project pair must fail");
        assert!(matches!(error, WinxError::WorkspaceBindingMismatch { .. }));
    }

    #[tokio::test]
    async fn strict_remote_binding_requires_workspace_root() {
        let workspace = tempfile::tempdir().expect("workspace");
        let service = WinxService::with_isolation(SessionIsolation::Strict);
        bind(&service, "thread", workspace.path()).await;
        let request = request(
            "ReadFiles",
            &serde_json::json!({"file_paths": [workspace.path()], "thread_id": "thread"}),
        );

        let error = service
            .validate_workspace_coherence(
                &request,
                &RequestScope::default(),
                HttpSessionAffinity::Thread,
                None,
            )
            .await
            .expect_err("missing binding must fail closed");
        assert!(matches!(error, WinxError::WorkspaceBindingRequired { .. }));
    }

    #[tokio::test]
    async fn strict_remote_initialize_accepts_a_file_as_the_workspace_hint() {
        let workspace = tempfile::tempdir().expect("workspace");
        let file = workspace.path().join("request.md");
        std::fs::write(&file, "request").expect("workspace file");
        let service = WinxService::with_isolation(SessionIsolation::Strict);
        let request = request(
            "Initialize",
            &serde_json::json!({
                "type": "first_call",
                "any_workspace_path": file,
                "mode_name": "wcgw",
                "thread_id": "thread"
            }),
        );

        let result = service
            .validate_workspace_coherence(
                &request,
                &RequestScope::default(),
                HttpSessionAffinity::Thread,
                None,
            )
            .await;
        assert_eq!(result.expect("file hint"), WorkspaceCoherence::FirstCall);
    }

    #[tokio::test]
    async fn strict_remote_initialize_recovers_safe_transitions_after_state_loss() {
        let workspace = tempfile::tempdir().expect("workspace");
        let service = WinxService::with_isolation(SessionIsolation::Strict);

        for (kind, thread_id) in
            [("reset_shell", "reset-thread"), ("user_asked_mode_change", "mode-thread")]
        {
            let request = request(
                "Initialize",
                &serde_json::json!({
                    "type": kind,
                    "any_workspace_path": workspace.path(),
                    "mode_name": "wcgw",
                    "thread_id": thread_id
                }),
            );
            let result = service
                .validate_workspace_coherence(
                    &request,
                    &RequestScope::default(),
                    HttpSessionAffinity::Thread,
                    None,
                )
                .await;
            assert_eq!(result.expect("missing session recovery"), WorkspaceCoherence::FirstCall);
        }
    }

    #[tokio::test]
    async fn strict_remote_initialize_keeps_workspace_change_fail_closed_after_state_loss() {
        let workspace = tempfile::tempdir().expect("workspace");
        let service = WinxService::with_isolation(SessionIsolation::Strict);
        let request = request(
            "Initialize",
            &serde_json::json!({
                "type": "user_asked_change_workspace",
                "any_workspace_path": workspace.path(),
                "mode_name": "wcgw",
                "thread_id": "workspace-change-thread"
            }),
        );

        let error = service
            .validate_workspace_coherence(
                &request,
                &RequestScope::default(),
                HttpSessionAffinity::Thread,
                None,
            )
            .await
            .expect_err("workspace changes must still require a separate session");
        assert!(matches!(error, WinxError::WorkspaceChangeRequiresNewSession { .. }));
    }
    #[tokio::test]
    async fn rejected_pair_preserves_binding_and_corrected_pair_still_validates() {
        let workspace = tempfile::tempdir().expect("bound workspace");
        let other = tempfile::tempdir().expect("other workspace");
        let workspace_root = canonical_workspace_identity(&workspace.path().to_string_lossy());

        for affinity in [HttpSessionAffinity::Thread, HttpSessionAffinity::Workspace] {
            let service = WinxService::with_isolation(SessionIsolation::Strict);
            let thread_id = expected_affinity_thread_id(&workspace_root, affinity, None)
                .unwrap_or_else(|| "preserved-thread".to_string());
            bind(&service, &thread_id, &workspace_root).await;
            let mismatched = request(
                "ReadFiles",
                &serde_json::json!({
                    "thread_id": &thread_id,
                    "workspace_root": other.path(),
                    "file_paths": [other.path().join("file.txt")]
                }),
            );
            let error = service
                .validate_workspace_coherence(&mismatched, &RequestScope::default(), affinity, None)
                .await
                .expect_err("wrong workspace must fail");
            match affinity {
                HttpSessionAffinity::Thread => {
                    assert!(matches!(error, WinxError::WorkspaceBindingMismatch { .. }));
                }
                HttpSessionAffinity::Workspace => {
                    assert!(matches!(error, WinxError::WorkspaceThreadMismatch { .. }));
                }
                HttpSessionAffinity::Conversation => {
                    unreachable!("test uses two explicit affinities")
                }
            }
            assert_eq!(service.bound_workspace(&thread_id).await, Some(workspace_root.clone()));

            let corrected = request(
                "ReadFiles",
                &serde_json::json!({
                    "thread_id": &thread_id,
                    "workspace_root": workspace.path(),
                    "file_paths": [workspace.path().join("file.txt")]
                }),
            );
            assert_eq!(
                service
                    .validate_workspace_coherence(
                        &corrected,
                        &RequestScope::default(),
                        affinity,
                        None,
                    )
                    .await
                    .expect("existing pair remains valid"),
                WorkspaceCoherence::Validated
            );
        }
    }
}