session.rs

  1use collections::{HashMap, HashSet};
  2use futures::{
  3    channel::{mpsc, oneshot},
  4    pin_mut, SinkExt, StreamExt,
  5};
  6use gpui::{AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity};
  7use mlua::{Lua, MultiValue, Table, UserData, UserDataMethods};
  8use parking_lot::Mutex;
  9use project::{search::SearchQuery, Fs, Project};
 10use regex::Regex;
 11use std::{
 12    cell::RefCell,
 13    path::{Path, PathBuf},
 14    sync::Arc,
 15};
 16use util::{paths::PathMatcher, ResultExt};
 17
 18use crate::{SCRIPT_END_TAG, SCRIPT_START_TAG};
 19
 20struct ForegroundFn(Box<dyn FnOnce(WeakEntity<ScriptSession>, AsyncApp) + Send>);
 21
 22pub struct ScriptSession {
 23    project: Entity<Project>,
 24    // TODO Remove this
 25    fs_changes: Arc<Mutex<HashMap<PathBuf, Vec<u8>>>>,
 26    foreground_fns_tx: mpsc::Sender<ForegroundFn>,
 27    _invoke_foreground_fns: Task<()>,
 28    scripts: Vec<Script>,
 29}
 30
 31impl ScriptSession {
 32    pub fn new(project: Entity<Project>, cx: &mut Context<Self>) -> Self {
 33        let (foreground_fns_tx, mut foreground_fns_rx) = mpsc::channel(128);
 34        ScriptSession {
 35            project,
 36            fs_changes: Arc::new(Mutex::new(HashMap::default())),
 37            foreground_fns_tx,
 38            _invoke_foreground_fns: cx.spawn(|this, cx| async move {
 39                while let Some(foreground_fn) = foreground_fns_rx.next().await {
 40                    foreground_fn.0(this.clone(), cx.clone());
 41                }
 42            }),
 43            scripts: Vec::new(),
 44        }
 45    }
 46
 47    pub fn new_script(&mut self) -> ScriptId {
 48        let id = ScriptId(self.scripts.len() as u32);
 49        let script = Script {
 50            id,
 51            state: ScriptState::Generating,
 52            source: SharedString::new_static(""),
 53        };
 54        self.scripts.push(script);
 55        id
 56    }
 57
 58    pub fn run_script(
 59        &mut self,
 60        script_id: ScriptId,
 61        script_src: String,
 62        cx: &mut Context<Self>,
 63    ) -> Task<anyhow::Result<()>> {
 64        let script = self.get_mut(script_id);
 65
 66        let stdout = Arc::new(Mutex::new(String::new()));
 67        script.source = script_src.clone().into();
 68        script.state = ScriptState::Running {
 69            stdout: stdout.clone(),
 70        };
 71
 72        let task = self.run_lua(script_src, stdout, cx);
 73
 74        cx.emit(ScriptEvent::Spawned(script_id));
 75
 76        cx.spawn(|session, mut cx| async move {
 77            let result = task.await;
 78
 79            session.update(&mut cx, |session, cx| {
 80                let script = session.get_mut(script_id);
 81                let stdout = script.stdout_snapshot();
 82
 83                script.state = match result {
 84                    Ok(()) => ScriptState::Succeeded { stdout },
 85                    Err(error) => ScriptState::Failed { stdout, error },
 86                };
 87
 88                cx.emit(ScriptEvent::Exited(script_id))
 89            })
 90        })
 91    }
 92
 93    fn run_lua(
 94        &mut self,
 95        script: String,
 96        stdout: Arc<Mutex<String>>,
 97        cx: &mut Context<Self>,
 98    ) -> Task<anyhow::Result<()>> {
 99        const SANDBOX_PREAMBLE: &str = include_str!("sandbox_preamble.lua");
100
101        // TODO Remove fs_changes
102        let fs_changes = self.fs_changes.clone();
103        // TODO Honor all worktrees instead of the first one
104        let root_dir = self
105            .project
106            .read(cx)
107            .visible_worktrees(cx)
108            .next()
109            .map(|worktree| worktree.read(cx).abs_path());
110
111        let fs = self.project.read(cx).fs().clone();
112        let foreground_fns_tx = self.foreground_fns_tx.clone();
113
114        let task = cx.background_spawn({
115            let stdout = stdout.clone();
116
117            async move {
118                let lua = Lua::new();
119                lua.set_memory_limit(2 * 1024 * 1024 * 1024)?; // 2 GB
120                let globals = lua.globals();
121                globals.set(
122                    "sb_print",
123                    lua.create_function({
124                        let stdout = stdout.clone();
125                        move |_, args: MultiValue| Self::print(args, &stdout)
126                    })?,
127                )?;
128                globals.set(
129                    "search",
130                    lua.create_async_function({
131                        let foreground_fns_tx = foreground_fns_tx.clone();
132                        let fs = fs.clone();
133                        move |lua, regex| {
134                            Self::search(lua, foreground_fns_tx.clone(), fs.clone(), regex)
135                        }
136                    })?,
137                )?;
138                globals.set(
139                    "sb_io_open",
140                    lua.create_function({
141                        let fs_changes = fs_changes.clone();
142                        let root_dir = root_dir.clone();
143                        move |lua, (path_str, mode)| {
144                            Self::io_open(&lua, &fs_changes, root_dir.as_ref(), path_str, mode)
145                        }
146                    })?,
147                )?;
148                globals.set("user_script", script)?;
149
150                lua.load(SANDBOX_PREAMBLE).exec_async().await?;
151
152                // Drop Lua instance to decrement reference count.
153                drop(lua);
154
155                anyhow::Ok(())
156            }
157        });
158
159        task
160    }
161
162    pub fn get(&self, script_id: ScriptId) -> &Script {
163        &self.scripts[script_id.0 as usize]
164    }
165
166    fn get_mut(&mut self, script_id: ScriptId) -> &mut Script {
167        &mut self.scripts[script_id.0 as usize]
168    }
169
170    /// Sandboxed print() function in Lua.
171    fn print(args: MultiValue, stdout: &Mutex<String>) -> mlua::Result<()> {
172        for (index, arg) in args.into_iter().enumerate() {
173            // Lua's `print()` prints tab characters between each argument.
174            if index > 0 {
175                stdout.lock().push('\t');
176            }
177
178            // If the argument's to_string() fails, have the whole function call fail.
179            stdout.lock().push_str(&arg.to_string()?);
180        }
181        stdout.lock().push('\n');
182
183        Ok(())
184    }
185
186    /// Sandboxed io.open() function in Lua.
187    fn io_open(
188        lua: &Lua,
189        fs_changes: &Arc<Mutex<HashMap<PathBuf, Vec<u8>>>>,
190        root_dir: Option<&Arc<Path>>,
191        path_str: String,
192        mode: Option<String>,
193    ) -> mlua::Result<(Option<Table>, String)> {
194        let root_dir = root_dir
195            .ok_or_else(|| mlua::Error::runtime("cannot open file without a root directory"))?;
196
197        let mode = mode.unwrap_or_else(|| "r".to_string());
198
199        // Parse the mode string to determine read/write permissions
200        let read_perm = mode.contains('r');
201        let write_perm = mode.contains('w') || mode.contains('a') || mode.contains('+');
202        let append = mode.contains('a');
203        let truncate = mode.contains('w');
204
205        // This will be the Lua value returned from the `open` function.
206        let file = lua.create_table()?;
207
208        // Store file metadata in the file
209        file.set("__path", path_str.clone())?;
210        file.set("__mode", mode.clone())?;
211        file.set("__read_perm", read_perm)?;
212        file.set("__write_perm", write_perm)?;
213
214        // Sandbox the path; it must be within root_dir
215        let path: PathBuf = {
216            let rust_path = Path::new(&path_str);
217
218            // Get absolute path
219            if rust_path.is_absolute() {
220                // Check if path starts with root_dir prefix without resolving symlinks
221                if !rust_path.starts_with(&root_dir) {
222                    return Ok((
223                        None,
224                        format!(
225                            "Error: Absolute path {} is outside the current working directory",
226                            path_str
227                        ),
228                    ));
229                }
230                rust_path.to_path_buf()
231            } else {
232                // Make relative path absolute relative to cwd
233                root_dir.join(rust_path)
234            }
235        };
236
237        // close method
238        let close_fn = {
239            let fs_changes = fs_changes.clone();
240            lua.create_function(move |_lua, file_userdata: mlua::Table| {
241                let write_perm = file_userdata.get::<bool>("__write_perm")?;
242                let path = file_userdata.get::<String>("__path")?;
243
244                if write_perm {
245                    // When closing a writable file, record the content
246                    let content = file_userdata.get::<mlua::AnyUserData>("__content")?;
247                    let content_ref = content.borrow::<FileContent>()?;
248                    let content_vec = content_ref.0.borrow();
249
250                    // Don't actually write to disk; instead, just update fs_changes.
251                    let path_buf = PathBuf::from(&path);
252                    fs_changes
253                        .lock()
254                        .insert(path_buf.clone(), content_vec.clone());
255                }
256
257                Ok(true)
258            })?
259        };
260        file.set("close", close_fn)?;
261
262        // If it's a directory, give it a custom read() and return early.
263        if path.is_dir() {
264            // TODO handle the case where we changed it in the in-memory fs
265
266            // Create a special directory handle
267            file.set("__is_directory", true)?;
268
269            // Store directory entries
270            let entries = match std::fs::read_dir(&path) {
271                Ok(entries) => {
272                    let mut entry_names = Vec::new();
273                    for entry in entries.flatten() {
274                        entry_names.push(entry.file_name().to_string_lossy().into_owned());
275                    }
276                    entry_names
277                }
278                Err(e) => return Ok((None, format!("Error reading directory: {}", e))),
279            };
280
281            // Save the list of entries
282            file.set("__dir_entries", entries)?;
283            file.set("__dir_position", 0usize)?;
284
285            // Create a directory-specific read function
286            let read_fn = lua.create_function(|_lua, file_userdata: mlua::Table| {
287                let position = file_userdata.get::<usize>("__dir_position")?;
288                let entries = file_userdata.get::<Vec<String>>("__dir_entries")?;
289
290                if position >= entries.len() {
291                    return Ok(None); // No more entries
292                }
293
294                let entry = entries[position].clone();
295                file_userdata.set("__dir_position", position + 1)?;
296
297                Ok(Some(entry))
298            })?;
299            file.set("read", read_fn)?;
300
301            // If we got this far, the directory was opened successfully
302            return Ok((Some(file), String::new()));
303        }
304
305        let fs_changes_map = fs_changes.lock();
306
307        let is_in_changes = fs_changes_map.contains_key(&path);
308        let file_exists = is_in_changes || path.exists();
309        let mut file_content = Vec::new();
310
311        if file_exists && !truncate {
312            if is_in_changes {
313                file_content = fs_changes_map.get(&path).unwrap().clone();
314            } else {
315                // Try to read existing content if file exists and we're not truncating
316                match std::fs::read(&path) {
317                    Ok(content) => file_content = content,
318                    Err(e) => return Ok((None, format!("Error reading file: {}", e))),
319                }
320            }
321        }
322
323        drop(fs_changes_map); // Unlock the fs_changes mutex.
324
325        // If in append mode, position should be at the end
326        let position = if append && file_exists {
327            file_content.len()
328        } else {
329            0
330        };
331        file.set("__position", position)?;
332        file.set(
333            "__content",
334            lua.create_userdata(FileContent(RefCell::new(file_content)))?,
335        )?;
336
337        // Create file methods
338
339        // read method
340        let read_fn = {
341            lua.create_function(
342                |_lua, (file_userdata, format): (mlua::Table, Option<mlua::Value>)| {
343                    let read_perm = file_userdata.get::<bool>("__read_perm")?;
344                    if !read_perm {
345                        return Err(mlua::Error::runtime("File not open for reading"));
346                    }
347
348                    let content = file_userdata.get::<mlua::AnyUserData>("__content")?;
349                    let mut position = file_userdata.get::<usize>("__position")?;
350                    let content_ref = content.borrow::<FileContent>()?;
351                    let content_vec = content_ref.0.borrow();
352
353                    if position >= content_vec.len() {
354                        return Ok(None); // EOF
355                    }
356
357                    match format {
358                        Some(mlua::Value::String(s)) => {
359                            let lossy_string = s.to_string_lossy();
360                            let format_str: &str = lossy_string.as_ref();
361
362                            // Only consider the first 2 bytes, since it's common to pass e.g. "*all"  instead of "*a"
363                            match &format_str[0..2] {
364                                "*a" => {
365                                    // Read entire file from current position
366                                    let result = String::from_utf8_lossy(&content_vec[position..])
367                                        .to_string();
368                                    position = content_vec.len();
369                                    file_userdata.set("__position", position)?;
370                                    Ok(Some(result))
371                                }
372                                "*l" => {
373                                    // Read next line
374                                    let mut line = Vec::new();
375                                    let mut found_newline = false;
376
377                                    while position < content_vec.len() {
378                                        let byte = content_vec[position];
379                                        position += 1;
380
381                                        if byte == b'\n' {
382                                            found_newline = true;
383                                            break;
384                                        }
385
386                                        // Skip \r in \r\n sequence but add it if it's alone
387                                        if byte == b'\r' {
388                                            if position < content_vec.len()
389                                                && content_vec[position] == b'\n'
390                                            {
391                                                position += 1;
392                                                found_newline = true;
393                                                break;
394                                            }
395                                        }
396
397                                        line.push(byte);
398                                    }
399
400                                    file_userdata.set("__position", position)?;
401
402                                    if !found_newline
403                                        && line.is_empty()
404                                        && position >= content_vec.len()
405                                    {
406                                        return Ok(None); // EOF
407                                    }
408
409                                    let result = String::from_utf8_lossy(&line).to_string();
410                                    Ok(Some(result))
411                                }
412                                "*n" => {
413                                    // Try to parse as a number (number of bytes to read)
414                                    match format_str.parse::<usize>() {
415                                        Ok(n) => {
416                                            let end =
417                                                std::cmp::min(position + n, content_vec.len());
418                                            let bytes = &content_vec[position..end];
419                                            let result = String::from_utf8_lossy(bytes).to_string();
420                                            position = end;
421                                            file_userdata.set("__position", position)?;
422                                            Ok(Some(result))
423                                        }
424                                        Err(_) => Err(mlua::Error::runtime(format!(
425                                            "Invalid format: {}",
426                                            format_str
427                                        ))),
428                                    }
429                                }
430                                "*L" => {
431                                    // Read next line keeping the end of line
432                                    let mut line = Vec::new();
433
434                                    while position < content_vec.len() {
435                                        let byte = content_vec[position];
436                                        position += 1;
437
438                                        line.push(byte);
439
440                                        if byte == b'\n' {
441                                            break;
442                                        }
443
444                                        // If we encounter a \r, add it and check if the next is \n
445                                        if byte == b'\r'
446                                            && position < content_vec.len()
447                                            && content_vec[position] == b'\n'
448                                        {
449                                            line.push(content_vec[position]);
450                                            position += 1;
451                                            break;
452                                        }
453                                    }
454
455                                    file_userdata.set("__position", position)?;
456
457                                    if line.is_empty() && position >= content_vec.len() {
458                                        return Ok(None); // EOF
459                                    }
460
461                                    let result = String::from_utf8_lossy(&line).to_string();
462                                    Ok(Some(result))
463                                }
464                                _ => Err(mlua::Error::runtime(format!(
465                                    "Unsupported format: {}",
466                                    format_str
467                                ))),
468                            }
469                        }
470                        Some(mlua::Value::Number(n)) => {
471                            // Read n bytes
472                            let n = n as usize;
473                            let end = std::cmp::min(position + n, content_vec.len());
474                            let bytes = &content_vec[position..end];
475                            let result = String::from_utf8_lossy(bytes).to_string();
476                            position = end;
477                            file_userdata.set("__position", position)?;
478                            Ok(Some(result))
479                        }
480                        Some(_) => Err(mlua::Error::runtime("Invalid format")),
481                        None => {
482                            // Default is to read a line
483                            let mut line = Vec::new();
484                            let mut found_newline = false;
485
486                            while position < content_vec.len() {
487                                let byte = content_vec[position];
488                                position += 1;
489
490                                if byte == b'\n' {
491                                    found_newline = true;
492                                    break;
493                                }
494
495                                // Handle \r\n
496                                if byte == b'\r' {
497                                    if position < content_vec.len()
498                                        && content_vec[position] == b'\n'
499                                    {
500                                        position += 1;
501                                        found_newline = true;
502                                        break;
503                                    }
504                                }
505
506                                line.push(byte);
507                            }
508
509                            file_userdata.set("__position", position)?;
510
511                            if !found_newline && line.is_empty() && position >= content_vec.len() {
512                                return Ok(None); // EOF
513                            }
514
515                            let result = String::from_utf8_lossy(&line).to_string();
516                            Ok(Some(result))
517                        }
518                    }
519                },
520            )?
521        };
522        file.set("read", read_fn)?;
523
524        // write method
525        let write_fn = {
526            let fs_changes = fs_changes.clone();
527
528            lua.create_function(move |_lua, (file_userdata, text): (mlua::Table, String)| {
529                let write_perm = file_userdata.get::<bool>("__write_perm")?;
530                if !write_perm {
531                    return Err(mlua::Error::runtime("File not open for writing"));
532                }
533
534                let content = file_userdata.get::<mlua::AnyUserData>("__content")?;
535                let position = file_userdata.get::<usize>("__position")?;
536                let content_ref = content.borrow::<FileContent>()?;
537                let mut content_vec = content_ref.0.borrow_mut();
538
539                let bytes = text.as_bytes();
540
541                // Ensure the vector has enough capacity
542                if position + bytes.len() > content_vec.len() {
543                    content_vec.resize(position + bytes.len(), 0);
544                }
545
546                // Write the bytes
547                for (i, &byte) in bytes.iter().enumerate() {
548                    content_vec[position + i] = byte;
549                }
550
551                // Update position
552                let new_position = position + bytes.len();
553                file_userdata.set("__position", new_position)?;
554
555                // Update fs_changes
556                let path = file_userdata.get::<String>("__path")?;
557                let path_buf = PathBuf::from(path);
558                fs_changes.lock().insert(path_buf, content_vec.clone());
559
560                Ok(true)
561            })?
562        };
563        file.set("write", write_fn)?;
564
565        // If we got this far, the file was opened successfully
566        Ok((Some(file), String::new()))
567    }
568
569    async fn search(
570        lua: Lua,
571        mut foreground_tx: mpsc::Sender<ForegroundFn>,
572        fs: Arc<dyn Fs>,
573        regex: String,
574    ) -> mlua::Result<Table> {
575        // TODO: Allow specification of these options.
576        let search_query = SearchQuery::regex(
577            &regex,
578            false,
579            false,
580            false,
581            PathMatcher::default(),
582            PathMatcher::default(),
583            None,
584        );
585        let search_query = match search_query {
586            Ok(query) => query,
587            Err(e) => return Err(mlua::Error::runtime(format!("Invalid search query: {}", e))),
588        };
589
590        // TODO: Should use `search_query.regex`. The tool description should also be updated,
591        // as it specifies standard regex.
592        let search_regex = match Regex::new(&regex) {
593            Ok(re) => re,
594            Err(e) => return Err(mlua::Error::runtime(format!("Invalid regex: {}", e))),
595        };
596
597        let mut abs_paths_rx =
598            Self::find_search_candidates(search_query, &mut foreground_tx).await?;
599
600        let mut search_results: Vec<Table> = Vec::new();
601        while let Some(path) = abs_paths_rx.next().await {
602            // Skip files larger than 1MB
603            if let Ok(Some(metadata)) = fs.metadata(&path).await {
604                if metadata.len > 1_000_000 {
605                    continue;
606                }
607            }
608
609            // Attempt to read the file as text
610            if let Ok(content) = fs.load(&path).await {
611                let mut matches = Vec::new();
612
613                // Find all regex matches in the content
614                for capture in search_regex.find_iter(&content) {
615                    matches.push(capture.as_str().to_string());
616                }
617
618                // If we found matches, create a result entry
619                if !matches.is_empty() {
620                    let result_entry = lua.create_table()?;
621                    result_entry.set("path", path.to_string_lossy().to_string())?;
622
623                    let matches_table = lua.create_table()?;
624                    for (ix, m) in matches.iter().enumerate() {
625                        matches_table.set(ix + 1, m.clone())?;
626                    }
627                    result_entry.set("matches", matches_table)?;
628
629                    search_results.push(result_entry);
630                }
631            }
632        }
633
634        // Create a table to hold our results
635        let results_table = lua.create_table()?;
636        for (ix, entry) in search_results.into_iter().enumerate() {
637            results_table.set(ix + 1, entry)?;
638        }
639
640        Ok(results_table)
641    }
642
643    async fn find_search_candidates(
644        search_query: SearchQuery,
645        foreground_tx: &mut mpsc::Sender<ForegroundFn>,
646    ) -> mlua::Result<mpsc::UnboundedReceiver<PathBuf>> {
647        Self::run_foreground_fn(
648            "finding search file candidates",
649            foreground_tx,
650            Box::new(move |session, mut cx| {
651                session.update(&mut cx, |session, cx| {
652                    session.project.update(cx, |project, cx| {
653                        project.worktree_store().update(cx, |worktree_store, cx| {
654                            // TODO: Better limit? For now this is the same as
655                            // MAX_SEARCH_RESULT_FILES.
656                            let limit = 5000;
657                            // TODO: Providing non-empty open_entries can make this a bit more
658                            // efficient as it can skip checking that these paths are textual.
659                            let open_entries = HashSet::default();
660                            let candidates = worktree_store.find_search_candidates(
661                                search_query,
662                                limit,
663                                open_entries,
664                                project.fs().clone(),
665                                cx,
666                            );
667                            let (abs_paths_tx, abs_paths_rx) = mpsc::unbounded();
668                            cx.spawn(|worktree_store, cx| async move {
669                                pin_mut!(candidates);
670
671                                while let Some(project_path) = candidates.next().await {
672                                    worktree_store.read_with(&cx, |worktree_store, cx| {
673                                        if let Some(worktree) = worktree_store
674                                            .worktree_for_id(project_path.worktree_id, cx)
675                                        {
676                                            if let Some(abs_path) = worktree
677                                                .read(cx)
678                                                .absolutize(&project_path.path)
679                                                .log_err()
680                                            {
681                                                abs_paths_tx.unbounded_send(abs_path)?;
682                                            }
683                                        }
684                                        anyhow::Ok(())
685                                    })??;
686                                }
687                                anyhow::Ok(())
688                            })
689                            .detach();
690                            abs_paths_rx
691                        })
692                    })
693                })
694            }),
695        )
696        .await
697    }
698
699    async fn run_foreground_fn<R: Send + 'static>(
700        description: &str,
701        foreground_tx: &mut mpsc::Sender<ForegroundFn>,
702        function: Box<dyn FnOnce(WeakEntity<Self>, AsyncApp) -> anyhow::Result<R> + Send>,
703    ) -> mlua::Result<R> {
704        let (response_tx, response_rx) = oneshot::channel();
705        let send_result = foreground_tx
706            .send(ForegroundFn(Box::new(move |this, cx| {
707                response_tx.send(function(this, cx)).ok();
708            })))
709            .await;
710        match send_result {
711            Ok(()) => (),
712            Err(err) => {
713                return Err(mlua::Error::runtime(format!(
714                    "Internal error while enqueuing work for {description}: {err}"
715                )))
716            }
717        }
718        match response_rx.await {
719            Ok(Ok(result)) => Ok(result),
720            Ok(Err(err)) => Err(mlua::Error::runtime(format!(
721                "Error while {description}: {err}"
722            ))),
723            Err(oneshot::Canceled) => Err(mlua::Error::runtime(format!(
724                "Internal error: response oneshot was canceled while {description}."
725            ))),
726        }
727    }
728}
729
730struct FileContent(RefCell<Vec<u8>>);
731
732impl UserData for FileContent {
733    fn add_methods<M: UserDataMethods<Self>>(_methods: &mut M) {
734        // FileContent doesn't have any methods so far.
735    }
736}
737
738#[derive(Debug)]
739pub enum ScriptEvent {
740    Spawned(ScriptId),
741    Exited(ScriptId),
742}
743
744impl EventEmitter<ScriptEvent> for ScriptSession {}
745
746#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
747pub struct ScriptId(u32);
748
749pub struct Script {
750    pub id: ScriptId,
751    pub state: ScriptState,
752    pub source: SharedString,
753}
754
755pub enum ScriptState {
756    Generating,
757    Running {
758        stdout: Arc<Mutex<String>>,
759    },
760    Succeeded {
761        stdout: String,
762    },
763    Failed {
764        stdout: String,
765        error: anyhow::Error,
766    },
767}
768
769impl Script {
770    pub fn source_tag(&self) -> String {
771        format!("{}{}{}", SCRIPT_START_TAG, self.source, SCRIPT_END_TAG)
772    }
773
774    /// If exited, returns a message with the output for the LLM
775    pub fn output_message_for_llm(&self) -> Option<String> {
776        match &self.state {
777            ScriptState::Generating { .. } => None,
778            ScriptState::Running { .. } => None,
779            ScriptState::Succeeded { stdout } => {
780                format!("Here's the script output:\n{}", stdout).into()
781            }
782            ScriptState::Failed { stdout, error } => format!(
783                "The script failed with:\n{}\n\nHere's the output it managed to print:\n{}",
784                error, stdout
785            )
786            .into(),
787        }
788    }
789
790    /// Get a snapshot of the script's stdout
791    pub fn stdout_snapshot(&self) -> String {
792        match &self.state {
793            ScriptState::Generating { .. } => String::new(),
794            ScriptState::Running { stdout } => stdout.lock().clone(),
795            ScriptState::Succeeded { stdout } => stdout.clone(),
796            ScriptState::Failed { stdout, .. } => stdout.clone(),
797        }
798    }
799
800    /// Returns the error if the script failed, otherwise None
801    pub fn error(&self) -> Option<&anyhow::Error> {
802        match &self.state {
803            ScriptState::Generating { .. } => None,
804            ScriptState::Running { .. } => None,
805            ScriptState::Succeeded { .. } => None,
806            ScriptState::Failed { error, .. } => Some(error),
807        }
808    }
809}
810
811#[cfg(test)]
812mod tests {
813    use gpui::TestAppContext;
814    use project::FakeFs;
815    use serde_json::json;
816    use settings::SettingsStore;
817
818    use super::*;
819
820    #[gpui::test]
821    async fn test_print(cx: &mut TestAppContext) {
822        let script = r#"
823            print("Hello", "world!")
824            print("Goodbye", "moon!")
825        "#;
826
827        let output = test_script(script, cx).await.unwrap();
828        assert_eq!(output, "Hello\tworld!\nGoodbye\tmoon!\n");
829    }
830
831    #[gpui::test]
832    async fn test_search(cx: &mut TestAppContext) {
833        let script = r#"
834            local results = search("world")
835            for i, result in ipairs(results) do
836                print("File: " .. result.path)
837                print("Matches:")
838                for j, match in ipairs(result.matches) do
839                    print("  " .. match)
840                end
841            end
842        "#;
843
844        let output = test_script(script, cx).await.unwrap();
845        assert_eq!(output, "File: /file1.txt\nMatches:\n  world\n");
846    }
847
848    async fn test_script(source: &str, cx: &mut TestAppContext) -> anyhow::Result<String> {
849        init_test(cx);
850        let fs = FakeFs::new(cx.executor());
851        fs.insert_tree(
852            "/",
853            json!({
854                "file1.txt": "Hello world!",
855                "file2.txt": "Goodbye moon!"
856            }),
857        )
858        .await;
859
860        let project = Project::test(fs, [Path::new("/")], cx).await;
861        let session = cx.new(|cx| ScriptSession::new(project, cx));
862
863        let (script_id, task) = session.update(cx, |session, cx| {
864            let script_id = session.new_script();
865            let task = session.run_script(script_id, source.to_string(), cx);
866
867            (script_id, task)
868        });
869
870        task.await?;
871
872        Ok(session.read_with(cx, |session, _cx| session.get(script_id).stdout_snapshot()))
873    }
874
875    fn init_test(cx: &mut TestAppContext) {
876        let settings_store = cx.update(SettingsStore::test);
877        cx.set_global(settings_store);
878        cx.update(Project::init_settings);
879    }
880}