feat(update): Suggest and accept qualified session names in update.

README.md

@@ -168,5 +168,7 @@ $ go work use . api git
     - cmd update.go no args
     - cmd update.go with args
     - cmd update.go completion suggestions
-  - Update should accept repo-qualified work unit names.
   - tmux hooks to automatically update session names when a session closes.
+  - Optimization ideas
+    - Add an existence filter to Sort for display_menu.
+    - Do blind Updates in update (will need to know whether anything changed).

cmd/update.go

@@ -32,48 +32,89 @@ var updateCommand = &cobra.Command{
 3. If given a work unit name, it will attempt to find that work unit in any of the repositories currently active in tmux and update both tmux and that VCS to point at the given work unit. Note: This means that you can update to a work unit that exists in a different repository.`,
 	Args: cobra.RangeArgs(0, 1),
 	ValidArgsFunction: func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
-		return suggestWorkUnitNames(cmd.Context(), toComplete), 0
+		return suggestWorkUnitNames(cmd.Context(), state.ParseSessionNameWithoutKnownRepository(toComplete)), 0
 	},
 	RunE: func(cmd *cobra.Command, args []string) error {
 		if len(args) == 0 {
 			return update(cmd.Context())
 		}
-		return updateTo(cmd.Context(), args[0])
+		return updateTo(cmd.Context(), state.ParseSessionNameWithoutKnownRepository(args[0]))
 	},
 }
 
-func suggestWorkUnitNames(ctx context.Context, toComplete string) []string {
-	vcs := api.Registered()
-	repos := make(map[state.RepoName]api.Repository)
-	if srv := tmux.MaybeCurrentServer(); srv != nil {
-		st, err := state.New(ctx, srv, vcs)
-		if err != nil {
-			slog.Warn("Could not determine repositories from tmux server.", "server", srv, "error", err)
-		} else {
-			repos = st.Repositories()
-		}
-	}
-	if repo, err := vcs.MaybeCurrentRepository(ctx); err != nil {
-		slog.Warn("Could not determine current repository.", "error", err)
-	} else {
-		n := state.NewRepoName(repo)
-		if _, ok := repos[n]; !ok {
-			repos[n] = repo
-		}
-	}
+func suggestWorkUnitNames(ctx context.Context, toComplete state.WorkUnitName) []string {
+	curRepo, repos := discoverRepositories(ctx, api.Registered())
 
 	var suggestions []string
 	for name, repo := range repos {
-		wus, err := repo.List(ctx, toComplete)
+		if toComplete.Repo != "" {
+			if toComplete.Repo != name.Repo {
+				// toComplete looks something like "foo>"
+				continue
+			}
+		} else if repo != curRepo && !strings.HasPrefix(name.Repo, toComplete.WorkUnit) {
+			// toComplete looks something like "foo"
+			continue
+		}
+
+		var prefix string
+		if toComplete.Repo != "" || repo == curRepo {
+			prefix = toComplete.WorkUnit
+		}
+
+		wus, err := repo.List(ctx, prefix)
 		if err != nil {
 			slog.Warn("Could not list work units.", "repo", name, "error", err)
+			continue
+		}
+		for _, wu := range wus {
+			if repo != curRepo {
+				wu = state.NewWorkUnitName(repo, wu).RepoString()
+			}
+			suggestions = append(suggestions, wu)
 		}
-		suggestions = append(suggestions, wus...)
 	}
-	suggestions = slices.DeleteFunc(suggestions, func(s string) bool { return !strings.HasPrefix(s, toComplete) })
+	suggestions = slices.DeleteFunc(suggestions, func(s string) bool { return !strings.HasPrefix(s, toComplete.String()) })
 	return suggestions
 }
 
+func discoverRepositories(ctx context.Context, vcs api.VersionControlSystems) (current api.Repository, all map[state.RepoName]api.Repository) {
+	all = make(map[state.RepoName]api.Repository)
+	var srv tmux.Server
+	var curSesh tmux.Session
+	if sesh := tmux.MaybeCurrentSession(); sesh != nil {
+		srv = sesh.Server()
+		curSesh = sesh
+	} else {
+		srv = tmux.DefaultServer()
+	}
+
+	st, err := state.New(ctx, srv, vcs)
+	if err != nil {
+		slog.Warn("Could not determine repositories from tmux server.", "server", srv, "error", err)
+	} else {
+		all = st.Repositories()
+		if curSesh != nil {
+			repo, _, err := st.WorkUnit(ctx, curSesh)
+			if err != nil {
+				slog.Warn("Could not determine current repository from tmux.", "server", srv, "error", err)
+			} else {
+				return repo, all
+			}
+		}
+	}
+
+	// If we're not in tmux or weren't able to discover the current repo from
+	// tmux, check for it directly.
+	cur, err := vcs.MaybeCurrentRepository(ctx)
+	if err != nil {
+		slog.Warn("Could not determine current repository.", "error", err)
+		return nil, all
+	}
+	all[state.NewRepoName(cur)] = cur
+	return cur, all
+}
+
 func update(ctx context.Context) error {
 	vcs := api.Registered()
 	curRepo, err := vcs.CurrentRepository(ctx)
@@ -129,7 +170,7 @@ func updateTmux(ctx context.Context, st *state.State, repo api.Repository, workU
 	return errors.Join(sesh.Server().AttachOrSwitch(ctx, sesh), err)
 }
 
-func updateTo(ctx context.Context, workUnitName string) error {
+func updateTo(ctx context.Context, sessionName state.WorkUnitName) error {
 	vcs := api.Registered()
 	srv := tmux.MaybeCurrentServer()
 	hasCurrentServer := srv != nil
@@ -141,32 +182,19 @@ func updateTo(ctx context.Context, workUnitName string) error {
 		return err
 	}
 
-	var repo api.Repository
-	if cur, err1 := existsInCurrentRepo(ctx, vcs, workUnitName); err1 == nil && cur != nil {
-		repo = cur
-	} else {
-		var err2 error
-		repo, err2 = st.MaybeFindRepository(ctx, workUnitName)
-		if err2 != nil {
-			return errors.Join(err1, err2)
-		}
-		if repo == nil {
-			return errors.Join(err1, fmt.Errorf("could not find any repository that contains work unit %q", workUnitName))
-		}
-		if err1 != nil {
-			slog.Warn("An error occurred with the current repository.", "error", err1)
-		}
+	repo, err := findRepository(ctx, vcs, st, sessionName)
+	if err != nil {
+		return err
 	}
-	slog.Info("Found repository for requested work unit.", "name", state.NewWorkUnitName(repo, workUnitName))
 
 	var update bool
 
 	// Update to the work unit.
 	if cur, err := repo.Current(ctx); err != nil {
 		return fmt.Errorf("couldn't check repo's current %s: %w", repo.VCS().WorkUnitName(), err)
-	} else if cur != workUnitName {
-		slog.Info("Updating repository.", "got", cur, "want", workUnitName)
-		if err := repo.Update(ctx, workUnitName); err != nil {
+	} else if cur != sessionName.WorkUnit {
+		slog.Info("Updating repository.", "got", cur, "want", sessionName.WorkUnit)
+		if err := repo.Update(ctx, sessionName.WorkUnit); err != nil {
 			return err
 		}
 		update = true
@@ -177,15 +205,15 @@ func updateTo(ctx context.Context, workUnitName string) error {
 	if !hasCurrentServer {
 		// Not currently attached to tmux.
 		needsSwitch = true
-	} else if sesh := st.Session(repo, workUnitName); sesh == nil {
+	} else if sesh := st.Session(repo, sessionName.WorkUnit); sesh == nil {
 		// Session doesn't exist.
 		needsSwitch = true
 	} else if cur := tmux.MaybeCurrentSession(); cur == nil || !tmux.SameSession(ctx, cur, sesh) {
 		// cur == nil shouldn't be possible. We already know we're attached to tmux.
 		needsSwitch = true
 	}
 	if needsSwitch {
-		if err := updateTmux(ctx, st, repo, workUnitName, !hasCurrentServer); err != nil {
+		if err := updateTmux(ctx, st, repo, sessionName.WorkUnit, !hasCurrentServer); err != nil {
 			return err
 		}
 		update = true
@@ -201,6 +229,28 @@ func updateTo(ctx context.Context, workUnitName string) error {
 	return nil
 }
 
+func findRepository(ctx context.Context, vcs api.VersionControlSystems, st *state.State, n state.WorkUnitName) (api.Repository, error) {
+	var err1, err2 error
+	if n.RepoName.Zero() {
+		cur, err1 := existsInCurrentRepo(ctx, vcs, n.WorkUnit)
+		if err1 == nil && cur != nil {
+			return cur, nil
+		}
+	}
+	repo, err2 := st.MaybeFindRepository(ctx, n)
+	if err2 != nil {
+		return nil, errors.Join(err1, err2)
+	}
+	if repo == nil {
+		return nil, errors.Join(err1, fmt.Errorf("could not find repository %v", n))
+	}
+	if err1 != nil {
+		slog.Warn("An error occurred with the current repository.", "error", err1)
+	}
+	slog.Info("Found repository for requested work unit.", "name", state.NewWorkUnitName(repo, n.WorkUnit))
+	return repo, nil
+}
+
 func existsInCurrentRepo(ctx context.Context, vcs api.VersionControlSystems, workUnitName string) (api.Repository, error) {
 	repo, err := vcs.MaybeCurrentRepository(ctx)
 	if err != nil {

tmux/state/state.go

@@ -144,7 +144,7 @@ func (st *State) WorkUnit(ctx context.Context, sesh tmux.Session) (api.Repositor
 	}
 	n, ok := st.sessionsByID[sesh.ID()]
 	if !ok {
-		return nil, "", fmt.Errorf("sesh does not have an associated work unit")
+		return nil, "", nil
 	}
 	return n.repo, n.workUnitName, nil
 }
@@ -299,18 +299,39 @@ func (st *State) updateSessionNames(ctx context.Context) error {
 // Returns an error if multiple api.Repositories claim that the given work unit
 // exists.
 // Returns nil, nil if no such api.Repository exists.
-func (st *State) MaybeFindRepository(ctx context.Context, workUnitName string) (api.Repository, error) {
-	repo, err := api.MaybeFindRepository(ctx, expmaps.Values(st.repos), func(repo api.Repository) (api.Repository, error) {
-		ok, err := repo.Exists(ctx, workUnitName)
-		if err != nil {
+func (st *State) MaybeFindRepository(ctx context.Context, n WorkUnitName) (api.Repository, error) {
+	var repos []api.Repository
+	switch {
+	case n.RepoName.VCS != "":
+		if n.RepoName.Repo == "" {
+			return nil, fmt.Errorf("WorkUnitName has VCS set, but not Repo: %v", n)
+		}
+
+		repo, ok := st.repos[n.RepoName]
+		if !ok {
+			return nil, nil
+		}
+		repos = append(repos, repo)
+	case n.RepoName.Repo != "":
+		for m, repo := range st.repos {
+			if n.Repo == m.Repo {
+				repos = append(repos, repo)
+			}
+		}
+	default:
+		repos = expmaps.Values(st.Repositories())
+	}
+
+	repo, err := api.MaybeFindRepository(ctx, repos, func(repo api.Repository) (api.Repository, error) {
+		if ok, err := repo.Exists(ctx, n.WorkUnit); err != nil {
 			return nil, err
-		} else if ok {
-			return repo, nil
+		} else if !ok {
+			return nil, nil
 		}
-		return nil, nil
+		return repo, nil
 	})
 	if err != nil {
-		return nil, fmt.Errorf("work unit %q: %w", workUnitName, err)
+		return nil, fmt.Errorf("work unit %v: %w", n, err)
 	}
 	return repo, nil
 }
@@ -323,6 +344,17 @@ func NewRepoName(repo api.Repository) RepoName {
 	return RepoName{VCS: repo.VCS().Name(), Repo: repo.Name()}
 }
 
+func (n RepoName) Zero() bool {
+	return n == RepoName{}
+}
+
+func (n RepoName) String() string {
+	if n.VCS != "" {
+		return fmt.Sprintf("%s>%s", n.VCS, n.Repo)
+	}
+	return n.Repo
+}
+
 func (n RepoName) LogValue() slog.Value {
 	return slog.GroupValue(slog.String("vcs", n.VCS), slog.String("repo", n.Repo))
 }
@@ -333,30 +365,45 @@ type WorkUnitName struct {
 }
 
 func ParseSessionName(repo api.Repository, tmuxSessionName string) WorkUnitName {
-	n := WorkUnitName{RepoName: NewRepoName(repo)}
+	n := ParseSessionNameWithoutKnownRepository(tmuxSessionName)
+	if m := NewRepoName(repo); n.RepoName != m {
+		if (n.RepoName.VCS != "" && n.RepoName.VCS != m.VCS) || (n.RepoName.Repo != "" && n.RepoName.Repo != m.Repo) {
+			slog.Warn("Session name does not agree with repository.", "session_name", tmuxSessionName, "repo", m)
+		}
+		n.RepoName = m
+	}
+	return n
+}
 
+func ParseSessionNameWithoutKnownRepository(tmuxSessionName string) WorkUnitName {
 	sp := strings.SplitN(tmuxSessionName, ">", 3)
 	switch len(sp) {
 	case 1:
-		n.WorkUnit = sp[0]
+		return WorkUnitName{WorkUnit: sp[0]}
 	case 2:
-		if n.Repo != sp[0] {
-			slog.Warn("Session name does not agree with repository.", "session_name", tmuxSessionName, "repo", n.Repo)
-		}
-		n.WorkUnit = sp[1]
+		return WorkUnitName{RepoName: RepoName{Repo: sp[0]}, WorkUnit: sp[1]}
 	default:
-		if n.VCS != sp[1] || n.Repo != sp[1] {
-			slog.Warn("Session name does not agree with repository.", "session_name", tmuxSessionName, "vcs", n.VCS, "repo", n.Repo)
-		}
-		n.WorkUnit = sp[2]
+		return WorkUnitName{RepoName: RepoName{VCS: sp[0], Repo: sp[1]}, WorkUnit: sp[2]}
 	}
-	return n
 }
 
 func NewWorkUnitName(repo api.Repository, workUnitName string) WorkUnitName {
 	return WorkUnitName{NewRepoName(repo), workUnitName}
 }
 
+func (n WorkUnitName) Zero() bool {
+	return n == WorkUnitName{}
+}
+
+func (n WorkUnitName) String() string {
+	if n.VCS != "" {
+		return fmt.Sprintf("%s>%s>%s", n.VCS, n.Repo, n.WorkUnit)
+	} else if n.Repo != "" {
+		return n.RepoString()
+	}
+	return n.WorkUnitString()
+}
+
 func (n WorkUnitName) RepoString() string {
 	return fmt.Sprintf("%s>%s", n.Repo, n.WorkUnit)
 }