Skip to main content

wisp/runtime/
dispatcher.rs

1use crate::command::{
2    Command, CommandResult, FilesystemCommand, GitCommand, GitWatchCommand, TerminalCommand,
3};
4use crate::git_review::{DiffScope, GitDiffEvent, GitWatchEvent};
5use crate::request::RequestId;
6use crate::runtime::{agent, files, git};
7use acp_utils::client::AcpClientHandle;
8use clankerdiff_git::GitRepository;
9use clankerdiff_watch::{RepositoryRequest, RepositoryState, RepositoryWatcher, WatchError, WatchOptions};
10use crossterm::{execute, style::Print};
11use futures::{StreamExt, future::poll_fn};
12use std::task::Poll;
13use tokio_stream::wrappers::WatchStream;
14use std::{io, path::PathBuf, sync::Arc};
15use tokio::sync::oneshot;
16
17use super::tasks::{ReadTask, TaskSupervisor};
18
19pub struct CommandDispatcher {
20    client_handle: AcpClientHandle,
21    tasks: TaskSupervisor,
22    git_review: Option<(RequestId, PathBuf)>,
23    git_repository: Option<GitRepository>,
24    git_watch: Option<RepositoryWatcher>,
25    git_updates: Option<WatchStream<RepositoryState>>,
26}
27
28impl CommandDispatcher {
29    pub fn new(client_handle: AcpClientHandle) -> Self {
30        Self {
31            client_handle,
32            tasks: TaskSupervisor::default(),
33            git_review: None,
34            git_repository: None,
35            git_watch: None,
36            git_updates: None,
37        }
38    }
39
40    pub fn dispatch(&mut self, command: Command) -> Option<CommandResult> {
41        match command {
42            Command::Agent(command) => agent::execute(&self.client_handle, command, &mut self.tasks),
43            Command::Filesystem(command) => {
44                let key = match &command {
45                    FilesystemCommand::IndexFiles { .. } => Some(ReadTask::FileIndex),
46                    FilesystemCommand::PrepareSubmission { .. } => Some(ReadTask::AttachmentPreparation),
47                    FilesystemCommand::ListThemes => Some(ReadTask::ThemeList),
48                    FilesystemCommand::ListReviewThemes => Some(ReadTask::ReviewThemeList),
49                    FilesystemCommand::ApplyTheme { .. } => None,
50                };
51                let work = async move { files::execute(command).await };
52                if let Some(key) = key {
53                    self.tasks.spawn_read(key, work);
54                } else {
55                    self.tasks.spawn_mutation(work);
56                }
57                None
58            }
59            Command::Git(GitCommand::Apply { review_id, action }) => {
60                let Some(repository) = self.git_repository.clone().filter(|_| self.is_review(review_id)) else {
61                    return Some(watch_stopped(review_id));
62                };
63                self.tasks.spawn_git_mutation(repository.root().to_path_buf(), async move {
64                    let result = repository.apply(action).await.map_err(Arc::new);
65                    CommandResult::GitDiff(GitDiffEvent { review_id, result })
66                });
67                None
68            }
69            Command::GitWatch(command) => {
70                match command {
71                    GitWatchCommand::Open { review_id, working_dir, scope } => {
72                        self.close_git_review();
73                        self.git_review = Some((review_id, working_dir));
74                        self.start_git_watch(scope);
75                    }
76                    GitWatchCommand::Refresh { review_id, scope } => {
77                        if !self.is_review(review_id) {
78                            return Some(watch_stopped(review_id));
79                        }
80                        if let Some(watcher) = &self.git_watch {
81                            let requests = watcher.request_tx.clone();
82                            self.tasks.spawn_read(ReadTask::GitRefresh, async move {
83                                    let (result_tx, completion) = oneshot::channel();
84                                    if requests.send(RepositoryRequest::SetScope { scope, result_tx }).await.is_err() {
85                                        return watch_stopped(review_id);
86                                    }
87                                    match completion.await {
88                                        Ok(result) => CommandResult::GitDiff(GitDiffEvent { review_id, result }),
89                                        Err(_) => watch_stopped(review_id),
90                                    }
91                            });
92                        } else {
93                            self.start_git_watch(scope);
94                        }
95                    }
96                    GitWatchCommand::Close { review_id } => {
97                        if self.is_review(review_id) {
98                            self.close_git_review();
99                        }
100                    }
101                }
102                None
103            }
104            Command::ResolveWorkspace { cwd } => {
105                self.tasks.spawn_read(ReadTask::Workspace, async move {
106                    let status = git::resolve_workspace_status(&cwd).await;
107                    CommandResult::WorkspaceResolved { cwd, status }
108                });
109                None
110            }
111            Command::Terminal(command) => execute_terminal(&command),
112        }
113    }
114
115    pub fn has_pending_tasks(&self) -> bool {
116        !self.tasks.is_empty() || self.git_updates.is_some()
117    }
118
119    pub async fn next_result(&mut self) -> Option<CommandResult> {
120        loop {
121            let result = poll_fn(|cx| {
122                if let (Some(updates), Some((review_id, _))) = (&mut self.git_updates, &self.git_review) {
123                    let review_id = *review_id;
124                    match updates.poll_next_unpin(cx) {
125                        Poll::Ready(Some(state)) => return Poll::Ready(Some(CommandResult::GitWatch(
126                            GitWatchEvent { review_id, result: Ok(state) },
127                        ))),
128                        Poll::Ready(None) => {
129                            self.close_git_review();
130                            return Poll::Ready(Some(watch_stopped(review_id)));
131                        }
132                        Poll::Pending => {}
133                    }
134                }
135                match self.tasks.poll_result(cx) {
136                    Poll::Ready(None) if self.git_updates.is_some() => Poll::Pending,
137                    result => result,
138                }
139            }).await?;
140            match result {
141                CommandResult::GitWatchStarted { review_id, result } if self.is_review(review_id) => {
142                    match result {
143                        Ok(started) => {
144                            let (repository, watcher) = *started;
145                            self.git_updates = Some(WatchStream::new(watcher.state_rx.clone()));
146                            self.git_repository = Some(repository);
147                            self.git_watch = Some(watcher);
148                            self.tasks.spawn_read(ReadTask::GitRefresh, async move {
149                                CommandResult::GitDiff(GitDiffEvent { review_id, result: Ok(()) })
150                            });
151                        }
152                        Err(error) => return Some(CommandResult::GitWatch(GitWatchEvent {
153                            review_id, result: Err(Arc::new(error)),
154                        })),
155                    }
156                }
157                CommandResult::GitWatchStarted { .. } => {}
158                result => return Some(result),
159            }
160        }
161    }
162
163    pub async fn shutdown(&mut self) {
164        self.close_git_review();
165        self.client_handle.disconnect().await;
166        self.tasks.shutdown().await;
167    }
168
169    fn is_review(&self, review_id: RequestId) -> bool {
170        self.git_review.as_ref().is_some_and(|(id, _)| *id == review_id)
171    }
172
173    fn start_git_watch(&mut self, scope: DiffScope) {
174        let Some((review_id, working_dir)) = self.git_review.clone() else { return; };
175        self.tasks.cancel_read(ReadTask::GitRefresh);
176        self.tasks.spawn_read(ReadTask::GitStart, async move {
177            let result = async {
178                let repository = GitRepository::discover(working_dir).await?;
179                let watcher = RepositoryWatcher::spawn(repository.clone(), scope, WatchOptions::default()).await?;
180                Ok(Box::new((repository, watcher)))
181            }.await;
182            CommandResult::GitWatchStarted { review_id, result }
183        });
184    }
185
186    fn close_git_review(&mut self) {
187        self.tasks.cancel_read(ReadTask::GitStart);
188        self.tasks.cancel_read(ReadTask::GitRefresh);
189        self.git_updates = None;
190        self.git_watch = None;
191        self.git_repository = None;
192        self.git_review = None;
193    }
194}
195
196fn watch_stopped(review_id: RequestId) -> CommandResult {
197    CommandResult::GitWatch(GitWatchEvent { review_id, result: Err(Arc::new(WatchError::Stopped)) })
198}
199
200fn execute_terminal(command: &TerminalCommand) -> Option<CommandResult> {
201    let TerminalCommand::RingBell = command;
202    execute!(io::stdout(), Print("\x07")).err().map(|error| CommandResult::TerminalFailed(error.to_string()))
203}