Add live git status indicators (#458)

crates/yaak-git/bindings/gen_git.ts

@@ -7,7 +7,11 @@ export type CloneResult = { "type": "success" } | { "type": "cancelled" } | { "t
 
 export type GitAuthor = { name: string | null, email: string | null, };
 
-export type GitCommit = { author: GitAuthor, when: string, message: string | null, };
+export type GitBranchInfo = { path: string, headRef: string | null, headRefShorthand: string | null, origins: Array<string>, localBranches: Array<string>, remoteBranches: Array<string>, ahead: number, behind: number, };
+
+export type GitCommit = { oid: string, author: GitAuthor, when: string, message: string | null, };
+
+export type GitFileDiff = { original: string, modified: string, };
 
 export type GitRemote = { name: string, url: string | null, };
 
@@ -17,6 +21,10 @@ export type GitStatusEntry = { relaPath: string, status: GitStatus, staged: bool
 
 export type GitStatusSummary = { path: string, headRef: string | null, headRefShorthand: string | null, entries: Array<GitStatusEntry>, origins: Array<string>, localBranches: Array<string>, remoteBranches: Array<string>, ahead: number, behind: number, };
 
+export type GitWorktreeStatus = { entries: Array<GitWorktreeStatusEntry>, };
+
+export type GitWorktreeStatusEntry = { relaPath: string, modelId: string | null, status: GitStatus, staged: boolean, };
+
 export type PullResult = { "type": "success", message: string, } | { "type": "up_to_date" } | { "type": "needs_credentials", url: string, error: string | null, } | { "type": "diverged", remote: string, branch: string, } | { "type": "uncommitted_changes" };
 
 export type PushResult = { "type": "success", message: string, } | { "type": "up_to_date" } | { "type": "needs_credentials", url: string, error: string | null, };

crates/yaak-git/index.ts

@@ -1,14 +1,18 @@
 import { useQuery } from "@tanstack/react-query";
-import { invoke } from "@tauri-apps/api/core";
+import { Channel, invoke } from "@tauri-apps/api/core";
+import { emit } from "@tauri-apps/api/event";
 import { createFastMutation } from "@yaakapp/yaak-client/hooks/useFastMutation";
 import { queryClient } from "@yaakapp/yaak-client/lib/queryClient";
 import { useMemo } from "react";
 import {
   BranchDeleteResult,
   CloneResult,
+  GitBranchInfo,
   GitCommit,
+  GitFileDiff,
   GitRemote,
   GitStatusSummary,
+  GitWorktreeStatus,
   PullResult,
   PushResult,
 } from "./bindings/gen_git";
@@ -26,6 +30,10 @@ export type DivergedStrategy = "force_reset" | "merge" | "cancel";
 
 export type UncommittedChangesStrategy = "reset" | "cancel";
 
+interface GitWatchResult {
+  unlistenEvent: string;
+}
+
 export interface GitCallbacks {
   addRemote: () => Promise<GitRemote | null>;
   promptCredentials: (
@@ -38,13 +46,98 @@ export interface GitCallbacks {
 
 const onSuccess = () => queryClient.invalidateQueries({ queryKey: ["git"] });
 
-export function useGit(dir: string, callbacks: GitCallbacks, refreshKey?: string) {
-  const mutations = useMemo(() => gitMutations(dir, callbacks), [dir, callbacks]);
-  const fetchAll = useQuery<void, string>({
+function gitWorktreeStatusQueryKey(dir?: string, refreshKey?: string) {
+  return refreshKey == null
+    ? (["git", "worktree_status", dir] as const)
+    : (["git", "worktree_status", dir, refreshKey] as const);
+}
+
+export function invalidateGitWorktreeStatus(dir?: string) {
+  return queryClient.invalidateQueries({ queryKey: gitWorktreeStatusQueryKey(dir) });
+}
+
+export function useGitWorktreeStatus(dir: string, refreshKey?: string) {
+  return useQuery<GitWorktreeStatus, string>({
+    queryKey: gitWorktreeStatusQueryKey(dir, refreshKey),
+    queryFn: () => invoke("cmd_git_worktree_status", { dir }),
+    placeholderData: (prev) => prev,
+  });
+}
+
+export function watchGitWorktreeStatus(dir: string, callback: (status: GitWorktreeStatus) => void) {
+  const channel = new Channel<GitWorktreeStatus>();
+  channel.onmessage = callback;
+  const unlistenPromise = invoke<GitWatchResult>("cmd_git_watch_worktree_status", {
+    dir,
+    channel,
+  });
+
+  void unlistenPromise
+    .then(({ unlistenEvent }) => {
+      addGitWatchKey(unlistenEvent);
+    })
+    .catch(console.debug);
+
+  return () =>
+    unlistenPromise
+      .then(async ({ unlistenEvent }) => {
+        unlistenGitWatcher(unlistenEvent);
+      })
+      .catch(console.error);
+}
+
+function useGitFetchAll(dir: string, refreshKey?: string) {
+  return useQuery<void, string>({
     queryKey: ["git", "fetch_all", dir, refreshKey],
     queryFn: () => invoke("cmd_git_fetch_all", { dir }),
     refetchInterval: 10 * 60_000,
   });
+}
+
+function useGitBranchInfoQuery(dir: string, refreshKey?: string, fetchAllUpdatedAt?: number) {
+  return useQuery<GitBranchInfo, string>({
+    refetchOnMount: true,
+    queryKey: ["git", "branch_info", dir, refreshKey, fetchAllUpdatedAt],
+    queryFn: () => invoke("cmd_git_branch_info", { dir }),
+    placeholderData: (prev) => prev,
+  });
+}
+
+export function useGitBranchInfo(dir: string, refreshKey?: string) {
+  const fetchAll = useGitFetchAll(dir, refreshKey);
+  return useGitBranchInfoQuery(dir, refreshKey, fetchAll.dataUpdatedAt);
+}
+
+export function useGitLog(dir: string, refreshKey?: string, relaPath?: string) {
+  return useQuery<GitCommit[], string>({
+    queryKey: ["git", "log", dir, refreshKey, relaPath],
+    queryFn: () =>
+      relaPath == null
+        ? invoke("cmd_git_log", { dir })
+        : invoke("cmd_git_log_for_file", { dir, relaPath }),
+    placeholderData: (prev) => prev,
+  });
+}
+
+export function useGitFileDiffForCommit(
+  dir: string,
+  relaPath: string,
+  commitOid: string | null | undefined,
+) {
+  return useQuery<GitFileDiff, string>({
+    enabled: commitOid != null,
+    queryKey: ["git", "file_diff_for_commit", dir, relaPath, commitOid],
+    queryFn: () => {
+      if (commitOid == null) throw new Error("Missing commit oid");
+      return invoke("cmd_git_file_diff_for_commit", { dir, relaPath, commitOid });
+    },
+  });
+}
+
+export function useGit(dir: string, callbacks: GitCallbacks, refreshKey?: string) {
+  const mutations = useGitMutations(dir, callbacks);
+  const fetchAll = useGitFetchAll(dir, refreshKey);
+
   return [
     {
       remotes: useQuery<GitRemote[], string>({
@@ -52,11 +145,7 @@ export function useGit(dir: string, callbacks: GitCallbacks, refreshKey?: string
         queryFn: () => getRemotes(dir),
         placeholderData: (prev) => prev,
       }),
-      log: useQuery<GitCommit[], string>({
-        queryKey: ["git", "log", dir, refreshKey],
-        queryFn: () => invoke("cmd_git_log", { dir }),
-        placeholderData: (prev) => prev,
-      }),
+      log: useGitLog(dir, refreshKey),
       status: useQuery<GitStatusSummary, string>({
         refetchOnMount: true,
         queryKey: ["git", "status", dir, refreshKey, fetchAll.dataUpdatedAt],
@@ -68,6 +157,10 @@ export function useGit(dir: string, callbacks: GitCallbacks, refreshKey?: string
   ] as const;
 }
 
+export function useGitMutations(dir: string, callbacks: GitCallbacks) {
+  return useMemo(() => gitMutations(dir, callbacks), [dir, callbacks]);
+}
+
 export const gitMutations = (dir: string, callbacks: GitCallbacks) => {
   const push = async () => {
     const remotes = await getRemotes(dir);
@@ -250,6 +343,20 @@ export const gitMutations = (dir: string, callbacks: GitCallbacks) => {
       mutationFn: () => invoke("cmd_git_reset_changes", { dir }),
       onSuccess,
     }),
+    restore: createFastMutation<void, string, { relaPaths: string[] }>({
+      mutationKey: ["git", "restore", dir],
+      mutationFn: (args) => invoke("cmd_git_restore_files", { dir, ...args }),
+      onSuccess,
+    }),
+    restoreFileFromCommit: createFastMutation<
+      void,
+      string,
+      { commitOid: string; relaPath: string }
+    >({
+      mutationKey: ["git", "restore-file-from-commit", dir],
+      mutationFn: (args) => invoke("cmd_git_restore_file_from_commit", { dir, ...args }),
+      onSuccess,
+    }),
   } as const;
 };
 
@@ -257,6 +364,35 @@ async function getRemotes(dir: string) {
   return invoke<GitRemote[]>("cmd_git_remotes", { dir });
 }
 
+function unlistenGitWatcher(unlistenEvent: string) {
+  void emit(unlistenEvent).then(() => {
+    removeGitWatchKey(unlistenEvent);
+  });
+}
+
+function getGitWatchKeys() {
+  return sessionStorage.getItem("git-worktree-watchers")?.split(",").filter(Boolean) ?? [];
+}
+
+function setGitWatchKeys(keys: string[]) {
+  sessionStorage.setItem("git-worktree-watchers", keys.join(","));
+}
+
+function addGitWatchKey(key: string) {
+  const keys = getGitWatchKeys();
+  setGitWatchKeys([...keys, key]);
+}
+
+function removeGitWatchKey(key: string) {
+  const keys = getGitWatchKeys();
+  setGitWatchKeys(keys.filter((k) => k !== key));
+}
+
+const gitWatchKeys = getGitWatchKeys();
+if (gitWatchKeys.length > 0) {
+  gitWatchKeys.forEach(unlistenGitWatcher);
+}
+
 /**
  * Clone a git repository, prompting for credentials if needed.
  */

crates/yaak-git/src/lib.rs

@@ -14,6 +14,7 @@ mod push;
 mod remotes;
 mod repository;
 mod reset;
+mod restore;
 mod status;
 mod unstage;
 mod util;
@@ -29,10 +30,15 @@ pub use commit::git_commit;
 pub use credential::git_add_credential;
 pub use fetch::git_fetch_all;
 pub use init::git_init;
-pub use log::{GitCommit, git_log};
+pub use log::{GitCommit, GitFileDiff, git_file_diff_for_commit, git_log, git_log_for_file};
 pub use pull::{PullResult, git_pull, git_pull_force_reset, git_pull_merge};
 pub use push::{PushResult, git_push};
 pub use remotes::{GitRemote, git_add_remote, git_remotes, git_rm_remote};
+pub use repository::{GitRepositoryPaths, git_path_is_ignored, git_repository_paths};
 pub use reset::git_reset_changes;
-pub use status::{GitStatusSummary, git_status};
+pub use restore::{git_restore, git_restore_file_from_commit};
+pub use status::{
+    GitBranchInfo, GitStatusSummary, GitWorktreeStatus, git_branch_info, git_status,
+    git_worktree_status,
+};
 pub use unstage::git_unstage;

crates/yaak-git/src/log.rs

@@ -8,6 +8,7 @@ use ts_rs::TS;
 #[serde(rename_all = "camelCase")]
 #[ts(export, export_to = "gen_git.ts")]
 pub struct GitCommit {
+    pub oid: String,
     pub author: GitAuthor,
     pub when: DateTime<Utc>,
     pub message: Option<String>,
@@ -21,7 +22,23 @@ pub struct GitAuthor {
     pub email: Option<String>,
 }
 
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
+#[serde(rename_all = "camelCase")]
+#[ts(export, export_to = "gen_git.ts")]
+pub struct GitFileDiff {
+    pub original: String,
+    pub modified: String,
+}
+
 pub fn git_log(dir: &Path) -> crate::error::Result<Vec<GitCommit>> {
+    git_log_inner(dir, None)
+}
+
+pub fn git_log_for_file(dir: &Path, rela_path: &Path) -> crate::error::Result<Vec<GitCommit>> {
+    git_log_inner(dir, Some(rela_path))
+}
+
+fn git_log_inner(dir: &Path, rela_path: Option<&Path>) -> crate::error::Result<Vec<GitCommit>> {
     let repo = open_repo(dir)?;
 
     // Return empty if empty repo or no head (new repo)
@@ -46,8 +63,16 @@ pub fn git_log(dir: &Path) -> crate::error::Result<Vec<GitCommit>> {
         .filter_map(|oid| {
             let oid = filter_try!(oid);
             let commit = filter_try!(repo.find_commit(oid));
+            if let Some(rela_path) = rela_path {
+                let touches_path = filter_try!(commit_touches_path(&repo, &commit, rela_path));
+                if !touches_path {
+                    return None;
+                }
+            }
+
             let author = commit.author();
             Some(GitCommit {
+                oid: oid.to_string(),
                 author: GitAuthor {
                     name: author.name().map(|s| s.to_string()),
                     email: author.email().map(|s| s.to_string()),
@@ -61,6 +86,53 @@ pub fn git_log(dir: &Path) -> crate::error::Result<Vec<GitCommit>> {
     Ok(log)
 }
 
+pub fn git_file_diff_for_commit(
+    dir: &Path,
+    commit_oid: &str,
+    rela_path: &Path,
+) -> crate::error::Result<GitFileDiff> {
+    let repo = open_repo(dir)?;
+    let oid = git2::Oid::from_str(commit_oid)?;
+    let commit = repo.find_commit(oid)?;
+    let new_tree = commit.tree()?;
+    let old_tree = if commit.parent_count() > 0 { Some(commit.parent(0)?.tree()?) } else { None };
+
+    Ok(GitFileDiff {
+        original: blob_text_at_path(&repo, old_tree.as_ref(), rela_path)?,
+        modified: blob_text_at_path(&repo, Some(&new_tree), rela_path)?,
+    })
+}
+
+fn commit_touches_path(
+    repo: &git2::Repository,
+    commit: &git2::Commit,
+    rela_path: &Path,
+) -> crate::error::Result<bool> {
+    let new_tree = commit.tree()?;
+    let old_tree = if commit.parent_count() > 0 { Some(commit.parent(0)?.tree()?) } else { None };
+
+    let mut opts = git2::DiffOptions::new();
+    opts.pathspec(rela_path);
+
+    let diff = repo.diff_tree_to_tree(old_tree.as_ref(), Some(&new_tree), Some(&mut opts))?;
+    Ok(diff.deltas().len() > 0)
+}
+
+fn blob_text_at_path(
+    repo: &git2::Repository,
+    tree: Option<&git2::Tree>,
+    rela_path: &Path,
+) -> crate::error::Result<String> {
+    let Some(tree) = tree else {
+        return Ok(String::new());
+    };
+    let Ok(entry) = tree.get_path(rela_path) else {
+        return Ok(String::new());
+    };
+    let blob = entry.to_object(repo)?.peel_to_blob()?;
+    Ok(String::from_utf8(blob.content().to_vec())?)
+}
+
 #[cfg(test)]
 fn convert_git_time_to_date(_git_time: git2::Time) -> DateTime<Utc> {
     DateTime::from_timestamp(0, 0).unwrap()

crates/yaak-git/src/repository.rs

@@ -1,5 +1,12 @@
 use crate::error::Error::{GitRepoNotFound, GitUnknown};
-use std::path::Path;
+use crate::error::{Error, Result};
+use std::path::{Path, PathBuf};
+
+#[derive(Debug, Clone)]
+pub struct GitRepositoryPaths {
+    pub workdir: PathBuf,
+    pub gitdir: PathBuf,
+}
 
 pub(crate) fn open_repo(dir: &Path) -> crate::error::Result<git2::Repository> {
     match git2::Repository::discover(dir) {
@@ -8,3 +15,17 @@ pub(crate) fn open_repo(dir: &Path) -> crate::error::Result<git2::Repository> {
         Err(e) => Err(GitUnknown(e)),
     }
 }
+
+pub fn git_repository_paths(dir: &Path) -> Result<GitRepositoryPaths> {
+    let repo = open_repo(dir)?;
+    let workdir = repo
+        .workdir()
+        .ok_or_else(|| Error::GenericError("Git repository does not have a worktree".into()))?
+        .to_path_buf();
+    Ok(GitRepositoryPaths { workdir, gitdir: repo.path().to_path_buf() })
+}
+
+pub fn git_path_is_ignored(dir: &Path, rela_path: &Path) -> Result<bool> {
+    let repo = open_repo(dir)?;
+    Ok(repo.status_should_ignore(rela_path)?)
+}

crates/yaak-git/src/restore.rs

@@ -0,0 +1,76 @@
+use crate::error::Result;
+use crate::repository::open_repo;
+use log::info;
+use std::fs;
+use std::path::{Component, Path};
+
+pub fn git_restore(dir: &Path, rela_path: &Path) -> Result<()> {
+    let repo = open_repo(dir)?;
+    validate_relative_path(rela_path)?;
+
+    let status = repo.status_file(rela_path).ok();
+    let is_untracked = status
+        .is_some_and(|s| s.contains(git2::Status::WT_NEW) || s.contains(git2::Status::INDEX_NEW));
+
+    info!("Restoring file {rela_path:?} in {dir:?}");
+    if is_untracked {
+        let mut index = repo.index()?;
+        let _ = index.remove_path(rela_path);
+        index.write()?;
+
+        let path = repo.workdir().unwrap_or(dir).join(rela_path);
+        if path.is_dir() {
+            fs::remove_dir_all(path)?;
+        } else if path.exists() {
+            fs::remove_file(path)?;
+        }
+        return Ok(());
+    }
+
+    let head = repo.head()?;
+    let commit = head.peel_to_commit()?;
+    repo.reset_default(Some(commit.as_object()), &[rela_path])?;
+
+    let mut checkout = git2::build::CheckoutBuilder::new();
+    checkout.force().path(rela_path);
+    repo.checkout_head(Some(&mut checkout))?;
+
+    Ok(())
+}
+
+pub fn git_restore_file_from_commit(dir: &Path, commit_oid: &str, rela_path: &Path) -> Result<()> {
+    let repo = open_repo(dir)?;
+    validate_relative_path(rela_path)?;
+
+    let oid = git2::Oid::from_str(commit_oid)?;
+    let commit = repo.find_commit(oid)?;
+    let tree = commit.tree()?;
+    let path = repo.workdir().unwrap_or(dir).join(rela_path);
+
+    info!("Restoring file {rela_path:?} from commit {commit_oid} in {dir:?}");
+    if tree.get_path(rela_path).is_err() {
+        if path.is_dir() {
+            fs::remove_dir_all(path)?;
+        } else if path.exists() {
+            fs::remove_file(path)?;
+        }
+        return Ok(());
+    }
+
+    let mut checkout = git2::build::CheckoutBuilder::new();
+    checkout.force().path(rela_path);
+    repo.checkout_tree(commit.as_object(), Some(&mut checkout))?;
+
+    Ok(())
+}
+
+fn validate_relative_path(path: &Path) -> Result<()> {
+    let is_safe = !path.as_os_str().is_empty()
+        && !path.is_absolute()
+        && path.components().all(|c| matches!(c, Component::Normal(_)));
+    if is_safe {
+        Ok(())
+    } else {
+        Err(crate::error::Error::GenericError(format!("Invalid restore path {}", path.display())))
+    }
+}

crates/yaak-git/src/status.rs

@@ -22,6 +22,20 @@ pub struct GitStatusSummary {
     pub behind: u32,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, TS, PartialEq)]
+#[serde(rename_all = "camelCase")]
+#[ts(export, export_to = "gen_git.ts")]
+pub struct GitBranchInfo {
+    pub path: String,
+    pub head_ref: Option<String>,
+    pub head_ref_shorthand: Option<String>,
+    pub origins: Vec<String>,
+    pub local_branches: Vec<String>,
+    pub remote_branches: Vec<String>,
+    pub ahead: u32,
+    pub behind: u32,
+}
+
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
 #[serde(rename_all = "camelCase")]
 #[ts(export, export_to = "gen_git.ts")]
@@ -33,6 +47,23 @@ pub struct GitStatusEntry {
     pub next: Option<SyncModel>,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, TS, PartialEq)]
+#[serde(rename_all = "camelCase")]
+#[ts(export, export_to = "gen_git.ts")]
+pub struct GitWorktreeStatus {
+    pub entries: Vec<GitWorktreeStatusEntry>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, TS, PartialEq)]
+#[serde(rename_all = "camelCase")]
+#[ts(export, export_to = "gen_git.ts")]
+pub struct GitWorktreeStatusEntry {
+    pub rela_path: String,
+    pub model_id: Option<String>,
+    pub status: GitStatus,
+    pub staged: bool,
+}
+
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, TS)]
 #[serde(rename_all = "snake_case")]
 #[ts(export, export_to = "gen_git.ts")]
@@ -46,31 +77,43 @@ pub enum GitStatus {
     TypeChange,
 }
 
+pub fn git_worktree_status(dir: &Path) -> crate::error::Result<GitWorktreeStatus> {
+    let repo = open_repo(dir)?;
+    let mut opts = git2::StatusOptions::new();
+    opts.include_ignored(false)
+        .include_untracked(true)
+        .recurse_untracked_dirs(true)
+        .include_unmodified(false);
+
+    let mut entries = Vec::new();
+    for entry in repo.statuses(Some(&mut opts))?.into_iter() {
+        let Some(rela_path) = entry.path() else {
+            continue;
+        };
+        let Some((status, staged)) = git_status_from_raw(entry.status()) else {
+            continue;
+        };
+
+        entries.push(GitWorktreeStatusEntry {
+            rela_path: rela_path.to_string(),
+            model_id: model_id_from_rela_path(Path::new(rela_path)),
+            status,
+            staged,
+        });
+    }
+
+    Ok(GitWorktreeStatus { entries })
+}
+
+pub fn git_branch_info(dir: &Path) -> crate::error::Result<GitBranchInfo> {
+    let repo = open_repo(dir)?;
+    git_branch_info_for_repo(&repo, dir)
+}
+
 pub fn git_status(dir: &Path) -> crate::error::Result<GitStatusSummary> {
     let repo = open_repo(dir)?;
-    let (head_tree, head_ref, head_ref_shorthand) = match repo.head() {
-        Ok(head) => {
-            let tree = head.peel_to_tree().ok();
-            let head_ref_shorthand = head.shorthand().map(|s| s.to_string());
-            let head_ref = head.name().map(|s| s.to_string());
-
-            (tree, head_ref, head_ref_shorthand)
-        }
-        Err(_) => {
-            // For "unborn" repos, reading from HEAD is the only way to get the branch name
-            // See https://github.com/starship/starship/pull/1336
-            let head_path = repo.path().join("HEAD");
-            let head_ref = fs::read_to_string(&head_path)
-                .ok()
-                .unwrap_or_default()
-                .lines()
-                .next()
-                .map(|s| s.trim_start_matches("ref:").trim().to_string());
-            let head_ref_shorthand =
-                head_ref.clone().map(|r| r.split('/').last().unwrap_or("unknown").to_string());
-            (None, head_ref, head_ref_shorthand)
-        }
-    };
+    let branch_info = git_branch_info_for_repo(&repo, dir)?;
+    let head_tree = repo.head().ok().and_then(|head| head.peel_to_tree().ok());
 
     let mut opts = git2::StatusOptions::new();
     opts.include_ignored(false)
@@ -83,51 +126,8 @@ pub fn git_status(dir: &Path) -> crate::error::Result<GitStatusSummary> {
     let mut entries: Vec<GitStatusEntry> = Vec::new();
     for entry in repo.statuses(Some(&mut opts))?.into_iter() {
         let rela_path = entry.path().unwrap().to_string();
-        let status = entry.status();
-        let index_status = match status {
-            // Note: order matters here, since we're checking a bitmap!
-            s if s.contains(git2::Status::CONFLICTED) => GitStatus::Conflict,
-            s if s.contains(git2::Status::INDEX_NEW) => GitStatus::Untracked,
-            s if s.contains(git2::Status::INDEX_MODIFIED) => GitStatus::Modified,
-            s if s.contains(git2::Status::INDEX_DELETED) => GitStatus::Removed,
-            s if s.contains(git2::Status::INDEX_RENAMED) => GitStatus::Renamed,
-            s if s.contains(git2::Status::INDEX_TYPECHANGE) => GitStatus::TypeChange,
-            s if s.contains(git2::Status::CURRENT) => GitStatus::Current,
-            s => {
-                warn!("Unknown index status {s:?}");
-                continue;
-            }
-        };
-
-        let worktree_status = match status {
-            // Note: order matters here, since we're checking a bitmap!
-            s if s.contains(git2::Status::CONFLICTED) => GitStatus::Conflict,
-            s if s.contains(git2::Status::WT_NEW) => GitStatus::Untracked,
-            s if s.contains(git2::Status::WT_MODIFIED) => GitStatus::Modified,
-            s if s.contains(git2::Status::WT_DELETED) => GitStatus::Removed,
-            s if s.contains(git2::Status::WT_RENAMED) => GitStatus::Renamed,
-            s if s.contains(git2::Status::WT_TYPECHANGE) => GitStatus::TypeChange,
-            s if s.contains(git2::Status::CURRENT) => GitStatus::Current,
-            s => {
-                warn!("Unknown worktree status {s:?}");
-                continue;
-            }
-        };
-
-        let status = if index_status == GitStatus::Current {
-            worktree_status.clone()
-        } else {
-            index_status.clone()
-        };
-
-        let staged = if index_status == GitStatus::Current && worktree_status == GitStatus::Current
-        {
-            // No change, so can't be added
-            false
-        } else if index_status != GitStatus::Current {
-            true
-        } else {
-            false
+        let Some((status, staged)) = git_status_from_raw(entry.status()) else {
+            continue;
         };
 
         // Get previous content from Git, if it's in there
@@ -158,9 +158,27 @@ pub fn git_status(dir: &Path) -> crate::error::Result<GitStatusSummary> {
         })
     }
 
+    Ok(GitStatusSummary {
+        entries,
+        path: branch_info.path,
+        head_ref: branch_info.head_ref,
+        head_ref_shorthand: branch_info.head_ref_shorthand,
+        origins: branch_info.origins,
+        local_branches: branch_info.local_branches,
+        remote_branches: branch_info.remote_branches,
+        ahead: branch_info.ahead,
+        behind: branch_info.behind,
+    })
+}
+
+fn git_branch_info_for_repo(
+    repo: &git2::Repository,
+    dir: &Path,
+) -> crate::error::Result<GitBranchInfo> {
+    let (head_ref, head_ref_shorthand) = git_head_refs(repo);
     let origins = repo.remotes()?.into_iter().filter_map(|o| Some(o?.to_string())).collect();
-    let local_branches = local_branch_names(&repo)?;
-    let remote_branches = remote_branch_names(&repo)?;
+    let local_branches = local_branch_names(repo)?;
+    let remote_branches = remote_branch_names(repo)?;
 
     // Compute ahead/behind relative to remote tracking branch
     let (ahead, behind) = (|| -> Option<(usize, usize)> {
@@ -174,15 +192,85 @@ pub fn git_status(dir: &Path) -> crate::error::Result<GitStatusSummary> {
     })()
     .unwrap_or((0, 0));
 
-    Ok(GitStatusSummary {
-        entries,
-        origins,
+    Ok(GitBranchInfo {
         path: dir.to_string_lossy().to_string(),
         head_ref,
         head_ref_shorthand,
+        origins,
         local_branches,
         remote_branches,
         ahead: ahead as u32,
         behind: behind as u32,
     })
 }
+
+fn git_head_refs(repo: &git2::Repository) -> (Option<String>, Option<String>) {
+    match repo.head() {
+        Ok(head) => {
+            let head_ref = head.name().map(|s| s.to_string());
+            let head_ref_shorthand = head.shorthand().map(|s| s.to_string());
+            (head_ref, head_ref_shorthand)
+        }
+        Err(_) => {
+            // For "unborn" repos, reading from HEAD is the only way to get the branch name
+            // See https://github.com/starship/starship/pull/1336
+            let head_path = repo.path().join("HEAD");
+            let head_ref = fs::read_to_string(&head_path)
+                .ok()
+                .unwrap_or_default()
+                .lines()
+                .next()
+                .map(|s| s.trim_start_matches("ref:").trim().to_string());
+            let head_ref_shorthand =
+                head_ref.clone().map(|r| r.split('/').last().unwrap_or("unknown").to_string());
+            (head_ref, head_ref_shorthand)
+        }
+    }
+}
+
+fn git_status_from_raw(status: git2::Status) -> Option<(GitStatus, bool)> {
+    let index_status = match status {
+        // Note: order matters here, since we're checking a bitmap!
+        s if s.contains(git2::Status::CONFLICTED) => GitStatus::Conflict,
+        s if s.contains(git2::Status::INDEX_NEW) => GitStatus::Untracked,
+        s if s.contains(git2::Status::INDEX_MODIFIED) => GitStatus::Modified,
+        s if s.contains(git2::Status::INDEX_DELETED) => GitStatus::Removed,
+        s if s.contains(git2::Status::INDEX_RENAMED) => GitStatus::Renamed,
+        s if s.contains(git2::Status::INDEX_TYPECHANGE) => GitStatus::TypeChange,
+        s if s.contains(git2::Status::CURRENT) => GitStatus::Current,
+        s => {
+            warn!("Unknown index status {s:?}");
+            return None;
+        }
+    };
+
+    let worktree_status = match status {
+        // Note: order matters here, since we're checking a bitmap!
+        s if s.contains(git2::Status::CONFLICTED) => GitStatus::Conflict,
+        s if s.contains(git2::Status::WT_NEW) => GitStatus::Untracked,
+        s if s.contains(git2::Status::WT_MODIFIED) => GitStatus::Modified,
+        s if s.contains(git2::Status::WT_DELETED) => GitStatus::Removed,
+        s if s.contains(git2::Status::WT_RENAMED) => GitStatus::Renamed,
+        s if s.contains(git2::Status::WT_TYPECHANGE) => GitStatus::TypeChange,
+        s if s.contains(git2::Status::CURRENT) => GitStatus::Current,
+        s => {
+            warn!("Unknown worktree status {s:?}");
+            return None;
+        }
+    };
+
+    let status =
+        if index_status == GitStatus::Current { worktree_status } else { index_status.clone() };
+    let staged = index_status != GitStatus::Current;
+
+    Some((status, staged))
+}
+
+fn model_id_from_rela_path(path: &Path) -> Option<String> {
+    let ext = path.extension()?.to_str()?;
+    if ext != "yaml" && ext != "yml" && ext != "json" {
+        return None;
+    }
+
+    path.file_stem()?.to_str()?.strip_prefix("yaak.").map(String::from)
+}