3885 lines · 163.1 KB · Rust
RawDownload
use std::collections::{HashMap, HashSet};use std::future::Future;use std::path::{Path, PathBuf};use std::sync::Arc; use ::ai::index::full_source_code_embedding::manager::{    CodebaseIndexManager, CodebaseIndexManagerEvent,    FragmentMetadataLookupError as LocalFragmentMetadataLookupError,};use ::ai::index::full_source_code_embedding::{    ContentHash, FragmentMetadata as LocalFragmentMetadata, NodeHash,};use ::ai::project_context::model::{ProjectContextModel, ProjectContextModelEvent};use remote_server::proto::OpenBufferSuccess;use repo_metadata::repositories::{DetectedRepositories, RepoDetectionSource};use repo_metadata::{RepoMetadataEvent, RepoMetadataModel, RepositoryIdentifier};use warp_core::channel::ChannelState;use warp_core::{SessionId, safe_error};use warp_files::{FileModel, FileModelEvent};use warp_util::content_version::ContentVersion;use warp_util::file::FileId;use warp_util::local_or_remote_path::LocalOrRemotePath;use warp_util::standardized_path::StandardizedPath;use warpui::r#async::{Spawnable, SpawnableOutput, SpawnedFutureHandle};use warpui::platform::TerminationMode;use warpui::{Entity, ModelContext, ModelHandle, SingletonEntity}; use super::codebase_index_status::{    codebase_index_status_to_proto, disabled_codebase_index_status,    not_enabled_codebase_index_status, queued_codebase_index_status,    unavailable_codebase_index_status,};use super::diff_state_tracker::{    DiffModelKey, DiffStateUpdate, RemoteDiffStateManager, SubscribeOutcome,};use super::proto::{    Abort, Authenticate, BranchInfo, BufferEdit, BufferUpdatedPush, ClientMessage, CloseBuffer,    CodebaseIndexLimits, CodebaseIndexStatus, CodebaseIndexStatusUpdated,    CodebaseIndexStatusesSnapshot, CodebaseResyncMode, DeleteFile, DeleteFileResponse,    DeleteFileSuccess, DiscardFilesError, DiscardFilesResponse, DiscardFilesSuccess,    DropCodebaseIndex, ErrorCode, ErrorResponse, FailedFileRead, FileContextProto,    FileOperationError, FragmentMetadata as ProtoFragmentMetadata,    FragmentMetadataLookupError as ProtoFragmentMetadataLookupError,    FragmentMetadataLookupErrorCode, GetBranchesError, GetBranchesResponse, GetBranchesSuccess,    GetDiffStateResponse, GetFragmentMetadataFromHash, GetFragmentMetadataFromHashResponse,    GetFragmentMetadataFromHashSuccess, GitCommitChainMode, GitCommitChainRequest,    GitCommitChainResponse, GitCommitChainSuccess, GitCreatePrRequest, GitCreatePrResponse,    GitGenerateCommitMessageRequest, GitGenerateCommitMessageResponse,    GitGetCommittedBranchFilesRequest, GitGetCommittedBranchFilesResponse,    GitGetCommittedBranchFilesSuccess, GitHubPrInfoPush, GitHubRepositoryInfoPush, GitOpDelta,    GitOpError, GitPushRequest, GitPushResponse, GitStatusPush, HomeSkillMetadata, IndexCodebase,    Initialize, InitializeResponse, MissingFragmentMetadata, NavigatedToDirectory,    NavigatedToDirectoryResponse, OpenBuffer, OpenBufferResponse, ReadFileContextResponse,    RemoteAgentContextSnapshot, RemoteContextFileProto, RemoteSkillProto, ResolveConflict,    ResolveConflictResponse, ResolveConflictSuccess, ResyncCodebase, RipgrepSearchRequest,    RunCommandError, RunCommandErrorCode, RunCommandRequest, RunCommandResponse, RunCommandSuccess,    SaveBuffer, SaveBufferResponse, SaveBufferSuccess, ServerMessage, SessionBootstrapped,    TextEdit, UpdateGitHubPrInfo, UpdateGitHubRepoInfo, UpdateGitStatus, UploadHandoffSnapshot,    WriteFile, WriteFileResponse, WriteFileSuccess, client_message, delete_file_response,    discard_files_response, get_diff_state_response, get_fragment_metadata_from_hash_response,    git_commit_chain_response, git_create_pr_response, git_generate_commit_message_response,    git_get_committed_branch_files_response, git_push_response, host_scoped_request, notification,    remote_skill_proto, resolve_conflict_response, run_command_response, save_buffer_response,    server_message, session_scoped_request, write_file_response,};use super::server_buffer_tracker::{PendingBufferRequestKind, ServerBufferTracker};use super::{diff_state_proto, ripgrep_search};use crate::code::global_buffer_model::{GlobalBufferModel, GlobalBufferModelEvent};use crate::code_review::diff_state::{CommitChainMode, DiffMode, FileStatusInfo};use crate::code_review::git_repo_model::{GitRepoModels, GitRepoStatusModel};use crate::code_review::github_repo_model::{GitHubRepoEvent, GitHubRepoModel};#[cfg(feature = "local_tty")]use crate::terminal::local_shell::LocalShellState;use crate::terminal::shell::ShellType; /// How long the daemon waits with no connections before exiting.pub const GRACE_PERIOD: std::time::Duration = std::time::Duration::from_secs(10 * 60); /// Server-side cap on the number of branches returned by `GetBranches`./// Prevents a client from forcing the daemon to enumerate an arbitrarily/// large ref list.const MAX_BRANCH_COUNT_CAP: usize = 500; /// Unique identifier for a connected proxy session in daemon mode.pub type ConnectionId = uuid::Uuid;use super::protocol::RequestId;use crate::ai::agent::FileLocations;use crate::ai::blocklist::handoff::snapshot::upload_result_to_proto;use crate::ai::blocklist::{ReadFileContextResult, read_local_file_context};use crate::ai::skills::{    BundledSkill, SkillManager, SkillManagerEvent, bundled_skill_snapshot_protos,};use crate::auth::auth_state::{AuthState, AuthStateProvider};use crate::code_review::git_actions;use crate::features::FeatureFlag;use crate::server::server_api::ServerApiProvider;use crate::terminal::model::session::command_executor::{    ExecuteCommandOptions, LocalCommandExecutor,};use crate::util::git; /// Resolves the global bundled resources directory populated by the install/// script (see [`remote_server::setup::remote_server_bundled_resources_dir`]),/// expanding the shell-form `~/` prefix against this process's home directory.////// This deliberately does not use `warp_core::paths::bundled_resources_dir`,/// whose macOS behavior resolves resources inside an app bundle. The global/// location is version-independent: the last install wins, and slight skew/// against this daemon's version is accepted.fn daemon_bundled_resources_dir() -> Option<PathBuf> {    let dir = remote_server::setup::remote_server_bundled_resources_dir();    let suffix = dir.strip_prefix("~/")?;    let dir = dirs::home_dir()?.join(suffix);    dir.is_dir().then_some(dir)}fn remote_agent_context_snapshot(    revision: u64,    bundled_skills: &[RemoteSkillProto],    ctx: &warpui::AppContext,) -> RemoteAgentContextSnapshot {    let home_dir = dirs::home_dir()        .map(|path| path.to_string_lossy().into_owned())        .unwrap_or_default();    let mut skills = bundled_skills.to_vec();    skills.extend(        SkillManager::as_ref(ctx)            .home_skills()            .map(|skill| RemoteSkillProto {                path: skill.path.display_path(),                content: skill.content.clone(),                source: Some(remote_skill_proto::Source::Home(HomeSkillMetadata {})),            }),    );    skills.sort_by(|a, b| a.path.cmp(&b.path));    let mut global_rules = ProjectContextModel::as_ref(ctx)        .global_rules()        .map(|rule| RemoteContextFileProto {            path: rule.path.display_path(),            content: rule.content,        })        .collect::<Vec<_>>();    global_rules.sort_by(|a, b| a.path.cmp(&b.path));    RemoteAgentContextSnapshot {        revision,        home_dir,        skills,        global_rules,    }} /// Outcome of dispatching a request-style `ClientMessage`.////// Notifications (fire-and-forget messages like `SessionBootstrapped` and/// `Abort`) do not produce a `HandlerOutcome`; they are dispatched inline in/// `handle_message` and return early.#[allow(clippy::large_enum_variant)]enum HandlerOutcome {    /// The response is ready synchronously — the caller sends it immediately.    Sync(server_message::Message),    /// The handler initiated async work whose response will be sent later.    ///    /// When the handle is `Some`, the caller inserts it into `in_progress`    /// so the request can be cancelled via `Abort`. Removal on    /// completion/abort is arranged by [`ServerModel::spawn_request_handler`].    ///    /// `None` is used for async work whose completion is delivered through    /// a separate event subscription and is not currently cancellable via    /// `Abort` (e.g. `FileModel` events for file writes and deletes, which    /// are tracked by `FileId` in `pending_file_ops` rather than by    /// `RequestId` in `in_progress`).    Async(Option<SpawnedFutureHandle>),} struct CodebaseIndexRequest {    repo_path: PathBuf,}struct CodebaseIndexRequestParams<'a> {    operation_name: &'a str,    repo_path: String,    auth_token: String,    auth_operation: &'a str,    path_kind: CodebaseIndexRequestPathKind,} #[derive(Clone, Copy)]enum CodebaseIndexRequestPathKind {    Canonicalized,    Requested,} /// Tracks an in-flight file write or delete so the async completion/// event can be correlated back to the originating client request.enum FileOpKind {    Write,    Delete,} struct PendingFileOp {    request_id: RequestId,    conn_id: ConnectionId,    kind: FileOpKind,} /// Manages pending file operations and ensures that the corresponding/// `FileModel` entry is always cleaned up when an operation completes/// or fails, preventing `FileState` leaks.struct PendingFileOps {    ops: HashMap<FileId, PendingFileOp>,} impl PendingFileOps {    fn new() -> Self {        Self {            ops: HashMap::new(),        }    }     /// Registers a file path with `FileModel`, sets the initial version,    /// and tracks the pending operation. Returns the `FileId` and    /// `ContentVersion` for the caller to initiate the actual I/O.    fn insert(        &mut self,        path: &Path,        request_id: RequestId,        conn_id: ConnectionId,        kind: FileOpKind,        ctx: &mut ModelContext<ServerModel>,    ) -> (FileId, ContentVersion) {        let file_model = FileModel::handle(ctx);        let file_id = file_model.update(ctx, |m, ctx| m.register_file_path(path, false, ctx));        let version = ContentVersion::new();        file_model.update(ctx, |m, _| m.set_version(file_id, version));        self.ops.insert(            file_id,            PendingFileOp {                request_id,                conn_id,                kind,            },        );        (file_id, version)    }     fn get(&self, file_id: &FileId) -> Option<&PendingFileOp> {        self.ops.get(file_id)    }     /// Removes a pending operation and unsubscribes the file from `FileModel`,    /// preventing the `FileState` entry from leaking.    fn remove(        &mut self,        file_id: FileId,        ctx: &mut ModelContext<ServerModel>,    ) -> Option<PendingFileOp> {        let op = self.ops.remove(&file_id)?;        FileModel::handle(ctx).update(ctx, |m, ctx| m.unsubscribe(file_id, ctx));        Some(op)    }} /// The top-level server-side orchestrator model.////// Receives `ClientMessage`s from connected proxy sessions and routes/// `ServerMessage` responses and push notifications back through each/// connection's dedicated sender channel.pub struct ServerModel {    /// Per-connection outbound channels, keyed by `ConnectionId`.    ///    /// The daemon can serve multiple proxy connections simultaneously — one    /// per SSH session / Warp tab connecting to this host.  Each entry maps    /// a connection's `Uuid` to the channel the connection task drains to    /// write `ServerMessage`s back to its proxy.    connection_senders: HashMap<ConnectionId, async_channel::Sender<ServerMessage>>,    /// Per-connection set of repo roots for which we've already sent a    /// snapshot in this connection's lifetime.    ///    /// Used to avoid sending duplicate snapshots on repeated    /// `NavigatedToDirectory` calls while the user `cd`s within the same repo.    snapshot_sent_roots_by_connection: HashMap<ConnectionId, HashSet<StandardizedPath>>,    /// Abort handle for the active grace timer, if any.    /// Calling `.abort()` cancels the timer before it fires.    grace_timer_cancel: Option<SpawnedFutureHandle>,    /// Tracks in-progress requests that can be cancelled via `Abort`.    /// Calling `.abort()` on the handle cancels the background future and    /// triggers its `on_abort` callback.    in_progress: HashMap<RequestId, SpawnedFutureHandle>,    /// Stable host identifier generated once at process startup.    /// Returned in every `InitializeResponse` so clients can deduplicate    /// host-scoped models.    host_id: String,    /// Bundled skill source entries detected and rendered on the daemon.    bundled_skills: Vec<RemoteSkillProto>,    /// Latest revisioned full replacement of all daemon-host Agent Mode context.    remote_agent_context_snapshot: RemoteAgentContextSnapshot,    /// Connections that have already received the current snapshot revision.    remote_agent_context_snapshot_sent: HashSet<ConnectionId>,    /// Per-session command executors created from `SessionBootstrapped` notifications.    executors: HashMap<SessionId, Arc<LocalCommandExecutor>>,    /// Tracks in-flight file write/delete operations and handles cleanup.    pending_file_ops: PendingFileOps,    /// Daemon-wide auth credentials and user identity.    auth_state: Arc<AuthState>,    /// Tracks open buffers, per-buffer connection sets, and pending async    /// buffer requests (OpenBuffer, SaveBuffer).    buffers: ServerBufferTracker,    /// Manages per-(repo, mode) diff state models and per-connection subscriptions.    diff_states: ModelHandle<RemoteDiffStateManager>,    /// In-flight host-scoped requests whose response may be delivered on    /// a different connection if the originating connection disconnects.    host_scoped_requests: HashMap<RequestId, ConnectionId>,    /// Per-repo local git status models tracked on the daemon, keyed by repo    /// path. Created when `NavigatedToDirectory` resolves a git root or a    /// client requests a snapshot; each is subscribed so watcher ticks    /// broadcast `GitStatusPush` to every connection.    git_status_models: HashMap<StandardizedPath, ModelHandle<GitRepoStatusModel>>,    /// Per-repo local GitHub-info models tracked on the daemon, keyed by repo    /// path. Created lazily on the first GitHub-info notification; each is    /// subscribed so `gh`-driven changes broadcast PR-info and repository-info    /// pushes to every connection.    github_repo_models: HashMap<StandardizedPath, ModelHandle<GitHubRepoModel>>,    /// Connections subscribed (via navigation) to each repo's git status,    /// keyed by repo path. A repo's git-status *and* GitHub-info models live    /// while this set is non-empty and are evicted once the last connection    /// unsubscribes (navigates away or disconnects). Mirrors    /// `RemoteDiffStateManager`'s per-key connection sets, keyed by repo only.    git_status_subscribers: HashMap<StandardizedPath, HashSet<ConnectionId>>,    /// Each connection's current git repo (a connection is in at most one repo    /// at a time), so a navigation can move its subscription and a disconnect    /// can drop it.    git_status_repo_by_conn: HashMap<ConnectionId, StandardizedPath>,} impl Entity for ServerModel {    type Event = ();} impl SingletonEntity for ServerModel {} impl ServerModel {    pub fn new(ctx: &mut ModelContext<Self>) -> Self {        let host_id = uuid::Uuid::new_v4().to_string();        log::info!(            "Daemon started: PID={}, host_id={}",            std::process::id(),            host_id        );        let bundled_skills = Vec::new();        let remote_agent_context_snapshot = remote_agent_context_snapshot(1, &bundled_skills, ctx);        let mut model = Self {            connection_senders: HashMap::new(),            snapshot_sent_roots_by_connection: HashMap::new(),            grace_timer_cancel: None,            in_progress: HashMap::new(),            host_id,            bundled_skills,            remote_agent_context_snapshot,            remote_agent_context_snapshot_sent: HashSet::new(),            executors: HashMap::new(),            pending_file_ops: PendingFileOps::new(),            auth_state: AuthStateProvider::as_ref(ctx).get().clone(),            buffers: ServerBufferTracker::new(),            diff_states: ctx.add_model(|_| RemoteDiffStateManager::new()),            host_scoped_requests: HashMap::new(),            git_status_models: HashMap::new(),            github_repo_models: HashMap::new(),            git_status_subscribers: HashMap::new(),            git_status_repo_by_conn: HashMap::new(),        };        // Subscribe to FileModel and RepoMetadataModel events        // file operation results and repo metadata pushes are forwarded to all        // connected proxy sessions.        {            let file_model = FileModel::handle(ctx);            ctx.subscribe_to_model(&file_model, |me, _, event, ctx| {                let file_id = event.file_id();                let Some(pending_kind) = me.pending_file_ops.get(&file_id).map(|op| &op.kind)                else {                    return; // Not a file op we're tracking.                };                let response_message = match (event, pending_kind) {                    (FileModelEvent::FileSaved { .. }, FileOpKind::Write) => {                        server_message::Message::WriteFileResponse(WriteFileResponse {                            result: Some(write_file_response::Result::Success(WriteFileSuccess {})),                        })                    }                    (FileModelEvent::FileSaved { .. }, FileOpKind::Delete) => {                        server_message::Message::DeleteFileResponse(DeleteFileResponse {                            result: Some(delete_file_response::Result::Success(                                DeleteFileSuccess {},                            )),                        })                    }                    (FileModelEvent::FailedToSave { error, .. }, FileOpKind::Write) => {                        server_message::Message::WriteFileResponse(WriteFileResponse {                            result: Some(write_file_response::Result::Error(FileOperationError {                                message: format!("{error}"),                            })),                        })                    }                    (FileModelEvent::FailedToSave { error, .. }, FileOpKind::Delete) => {                        server_message::Message::DeleteFileResponse(DeleteFileResponse {                            result: Some(delete_file_response::Result::Error(FileOperationError {                                message: format!("{error}"),                            })),                        })                    }                    (FileModelEvent::FileLoaded { .. }, _)                    | (FileModelEvent::FailedToLoad { .. }, _)                    | (FileModelEvent::FileUpdated { .. }, _) => return,                };                // Remove the pending op and unsubscribe from FileModel.                let pending = me                    .pending_file_ops                    .remove(file_id, ctx)                    .expect("pending op was confirmed present");                me.send_server_message(                    Some(pending.conn_id),                    Some(&pending.request_id),                    response_message,                );            });        }        {            let repo_model = RepoMetadataModel::handle(ctx);            ctx.subscribe_to_model(&repo_model, |me, _, event, ctx| match event {                RepoMetadataEvent::IncrementalUpdateReady { update } => {                    me.send_server_message(                        None,                        None,                        server_message::Message::RepoMetadataUpdate(update.into()),                    );                }                RepoMetadataEvent::RepositoryUpdated {                    id: RepositoryIdentifier::Local(path),                } => {                    // A repo finished indexing — push the full tree as a snapshot.                    let id = RepositoryIdentifier::local(path.clone());                    let repo_model = RepoMetadataModel::handle(ctx);                    if let Some(state) = repo_model.as_ref(ctx).get_repository(&id, ctx) {                        let entries = super::repo_metadata_proto::file_tree_entry_to_snapshot_proto(                            &state.entry,                        );                        let standing_results = repo_model                            .as_ref(ctx)                            .standing_query_results(&id, ctx)                            .map(|results| (&results.as_snapshot_delta()).into());                        me.send_server_message(                            None,                            None,                            server_message::Message::RepoMetadataSnapshot(                                super::proto::RepoMetadataSnapshot {                                    repo_path: path.to_string(),                                    entries,                                    sync_complete: true,                                    standing_results,                                },                            ),                        );                        // Mark this root as snapshot-sent for all active connections                        // so subsequent NavigatedToDirectory calls skip re-sending.                        for sent_roots in me.snapshot_sent_roots_by_connection.values_mut() {                            sent_roots.insert(path.clone());                        }                    }                }                RepoMetadataEvent::RepositoryRemoved { .. }                | RepoMetadataEvent::FileTreeUpdated { .. }                | RepoMetadataEvent::FileTreeEntryUpdated { .. }                | RepoMetadataEvent::StandingQueryResultsUpdated { .. }                | RepoMetadataEvent::UpdatingRepositoryFailed { .. }                | RepoMetadataEvent::RepositoryUpdated {                    id: RepositoryIdentifier::Remote(_),                } => {}            });        }        let index_manager = CodebaseIndexManager::handle(ctx);        ctx.subscribe_to_model(&index_manager, |me, _, event, ctx| {            me.handle_codebase_index_manager_event(event, ctx);        });        // Subscribe to GlobalBufferModel events for server-local buffers.        {            let gbm = GlobalBufferModel::handle(ctx);            ctx.subscribe_to_model(&gbm, |me, _, event, ctx| match event {                GlobalBufferModelEvent::BufferLoaded { file_id, .. } => {                    // Complete all pending OpenBuffer requests for this file.                    let pending = me.buffers.take_pending_by_kind(                        file_id,                        PendingBufferRequestKind::OpenBuffer,                    );                    if !pending.is_empty() {                        let gbm = GlobalBufferModel::handle(ctx);                        let content = gbm.as_ref(ctx).content_for_file(*file_id, ctx);                        let server_version = gbm                            .as_ref(ctx)                            .sync_clock_for_server_local(*file_id)                            .map(|c| c.server_version.as_u64());                         for req in pending {                            let message = match (&content, server_version) {                                (Some(content), Some(sv)) => {

Lines 1–500 of 3885.

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t