Coverage for src / ai_shell / interactive.py: 99%

1"""Interactive multi-pane wizard for ``ai-shell claude --multi -i``. 

2 

3Walks the user through a guided menu to configure each tmux pane, 

4then converts the collected choices into :class:`~ai_shell.tmux.PaneSpec` 

5objects ready for :func:`~ai_shell.tmux.build_tmux_commands`. 

6""" 

7 

8from __future__ import annotations 

9 

10import enum 

11import uuid 

12from collections.abc import Callable 

13from dataclasses import dataclass, field 

14from typing import TYPE_CHECKING, Any 

15 

16import click 

17from rich.console import Console 

18 

19if TYPE_CHECKING: 

20 from ai_shell.tmux import PaneSpec 

21 

22 

23# ── Data model ─────────────────────────────────────────────────────── 

24 

25 

26class PaneType(enum.Enum): 

27 """Types of panes available in interactive mode.""" 

28 

29 THIS_PROJECT = "this_project" 

30 BASH = "bash" 

31 WORKSPACE_REPO = "workspace_repo" 

32 

33 

34@dataclass 

35class PaneChoice: 

36 """A user's selection for a single window.""" 

37 

38 pane_type: PaneType 

39 repo_name: str = "" 

40 repo_path: str = "" 

41 

42 

43@dataclass 

44class InteractiveConfig: 

45 """Collected results from the interactive wizard.""" 

46 

47 pane_choices: list[PaneChoice] = field(default_factory=list) 

48 team_mode: bool = False 

49 shared_chrome: bool = False 

50 

51 

52# ── Option builder ─────────────────────────────────────────────────── 

53 

54 

55@dataclass 

56class _PaneOption: 

57 """A single numbered option in the per-window menu.""" 

58 

59 label: str 

60 pane_type: PaneType 

61 repo_name: str = "" 

62 repo_path: str = "" 

63 

64 

65def _build_pane_options( 

66 project_name: str, 

67 workspace_repos: list[dict[str, Any]] | None, 

68) -> list[_PaneOption]: 

69 """Build the numbered option list for per-window type selection.""" 

70 options: list[_PaneOption] = [ 

71 _PaneOption( 

72 label=f"This project ({project_name}) - Claude in worktree", 

73 pane_type=PaneType.THIS_PROJECT, 

74 ), 

75 _PaneOption( 

76 label="Bash shell", 

77 pane_type=PaneType.BASH, 

78 ), 

79 ] 

80 if workspace_repos: 

81 for repo in workspace_repos: 

82 name = repo["name"] 

83 path = repo.get("path", f"./{name}") 

84 repo_type = repo.get("repo_type", "") 

85 label = f"{name} ({repo_type})" if repo_type else name 

86 options.append( 

87 _PaneOption( 

88 label=label, 

89 pane_type=PaneType.WORKSPACE_REPO, 

90 repo_name=name, 

91 repo_path=path, 

92 ) 

93 ) 

94 return options 

95 

96 

97# ── Wizard ─────────────────────────────────────────────────────────── 

98 

99 

100def run_interactive_wizard( 

101 *, 

102 project_name: str, 

103 workspace_repos: list[dict[str, Any]] | None = None, 

104 console: Console | None = None, 

105) -> InteractiveConfig | None: 

106 """Walk the user through the interactive multi-pane setup. 

107 

108 Returns :class:`InteractiveConfig` with all choices, or ``None`` if the 

109 user cancels (Ctrl-C / EOFError). 

110 """ 

111 if console is None: 

112 console = Console(stderr=True) 

113 

114 options = _build_pane_options(project_name, workspace_repos) 

115 

116 try: 

117 # Step 1: number of windows 

118 num_windows: int = click.prompt( 

119 "How many windows?", 

120 type=click.IntRange(2, 4), 

121 default=2, 

122 ) 

123 

124 # Step 2: per-window type 

125 choices: list[PaneChoice] = [] 

126 for win_idx in range(1, num_windows + 1): 

127 console.print() 

128 console.print(f" [bold]Window {win_idx}:[/bold]") 

129 for i, opt in enumerate(options, 1): 

130 console.print(f" {i}. {opt.label}") 

131 

132 selection: int = click.prompt( 

133 f" Select type for window {win_idx}", 

134 type=click.IntRange(1, len(options)), 

135 default=1, 

136 ) 

137 chosen = options[selection - 1] 

138 choices.append( 

139 PaneChoice( 

140 pane_type=chosen.pane_type, 

141 repo_name=chosen.repo_name, 

142 repo_path=chosen.repo_path, 

143 ) 

144 ) 

145 

146 # Step 3: pre-launch options 

147 console.print() 

148 team_mode = click.confirm( 

149 "Enable teams mode on the primary Claude pane?", 

150 default=False, 

151 ) 

152 shared_chrome = click.confirm( 

153 "Enable shared Chrome browser for all Claude panes?", 

154 default=False, 

155 ) 

156 

157 except (EOFError, KeyboardInterrupt, click.Abort): 

158 return None 

159 

160 return InteractiveConfig( 

161 pane_choices=choices, 

162 team_mode=team_mode, 

163 shared_chrome=shared_chrome, 

164 ) 

165 

166 

167# ── Pane builder ───────────────────────────────────────────────────── 

168 

169 

170def build_interactive_panes( 

171 *, 

172 config: InteractiveConfig, 

173 project_name: str, 

174 container_name: str, 

175 container_project_root: str, 

176 safe: bool, 

177 extra_args: tuple[str, ...], 

178 mcp_config_path: str | None = None, 

179 setup_worktree_fn: Callable[[str, str, str], str], 

180) -> list[PaneSpec]: 

181 """Convert :class:`InteractiveConfig` into pane specs for tmux. 

182 

183 Parameters 

184 ---------- 

185 setup_worktree_fn 

186 ``(container_name, project_dir, worktree_name) -> worktree_abs_path``. 

187 In production this is ``tools._setup_worktree``; in tests, a mock. 

188 """ 

189 from ai_shell.tmux import PaneSpec, build_claude_pane_command 

190 

191 panes: list[PaneSpec] = [] 

192 team_assigned = False 

193 bash_counter = 0 

194 

195 for choice in config.pane_choices: 

196 if choice.pane_type == PaneType.THIS_PROJECT: 

197 wt_name = uuid.uuid4().hex[:8] 

198 worktree_dir = setup_worktree_fn(container_name, container_project_root, wt_name) 

199 use_team = config.team_mode and not team_assigned 

200 pane_cmd = build_claude_pane_command( 

201 repo_name=project_name, 

202 safe=safe, 

203 extra_args=extra_args, 

204 worktree_name=wt_name, 

205 mcp_config_path=mcp_config_path if config.shared_chrome else None, 

206 team_env=use_team, 

207 ) 

208 if use_team: 

209 team_assigned = True 

210 panes.append( 

211 PaneSpec( 

212 name=f"{project_name}-wt-{wt_name}", 

213 command=pane_cmd, 

214 working_dir=worktree_dir, 

215 ) 

216 ) 

217 

218 elif choice.pane_type == PaneType.BASH: 

219 bash_counter += 1 

220 panes.append( 

221 PaneSpec( 

222 name=f"bash-{bash_counter}", 

223 command="/bin/bash", 

224 working_dir=container_project_root, 

225 ) 

226 ) 

227 

228 elif choice.pane_type == PaneType.WORKSPACE_REPO: 

229 rel = choice.repo_path.lstrip("./") 

230 working_dir = f"{container_project_root}/{rel}" 

231 use_team = config.team_mode and not team_assigned 

232 pane_cmd = build_claude_pane_command( 

233 repo_name=choice.repo_name, 

234 safe=safe, 

235 extra_args=extra_args, 

236 mcp_config_path=mcp_config_path if config.shared_chrome else None, 

237 team_env=use_team, 

238 ) 

239 if use_team: 

240 team_assigned = True 

241 panes.append( 

242 PaneSpec( 

243 name=choice.repo_name, 

244 command=pane_cmd, 

245 working_dir=working_dir, 

246 ) 

247 ) 

248 

249 return panes