main
1{ pkgs }:
2
3let
4 # Helper script to run a command in a specific directory with proper environment
5 # Usage: run-in-dir <command> <directory> [model]
6 # If model is provided, adds --model <model> to the command
7 run-in-dir = pkgs.writeShellScript "run-in-dir" ''
8 CMD="''${1:-}"
9 DIR="''${2:-}"
10 MODEL="''${3:-}"
11
12 # Set TERM for proper color support
13 export TERM=xterm-256color
14
15 # Command is required
16 if [ -z "$CMD" ]; then
17 echo "Error: command argument is required" >&2
18 exit 1
19 fi
20
21 # Directory is required
22 if [ -z "$DIR" ]; then
23 echo "Error: directory argument is required" >&2
24 exit 1
25 fi
26
27 # Check if directory exists
28 if [ ! -d "$DIR" ]; then
29 echo "Error: directory does not exist: $DIR" >&2
30 exit 1
31 fi
32
33 # Change to specified directory
34 cd "$DIR"
35
36 # Build command with optional model flag
37 if [ -n "$MODEL" ]; then
38 FULL_CMD="$CMD --model $MODEL"
39 else
40 FULL_CMD="$CMD"
41 fi
42
43 # Execute through zsh with proper environment
44 # zshenv sets up PATH and other environment variables
45 exec ${pkgs.zsh}/bin/zsh -c "source ~/.zshenv; exec $FULL_CMD"
46 '';
47
48in
49
50pkgs.writeScriptBin "shpool-ssh-wrapper" ''
51 #!${pkgs.bash}/bin/bash
52 # Shpool SSH wrapper - automatically runs commands for specific session patterns
53 #
54 # Syntax: session-name=directory (e.g., "claude/home=src/home")
55 #
56 # Supported session patterns:
57 # claude/* - Claude Code with Vertex AI (cr command)
58 # copilot/* - GitHub Copilot CLI
59 # opencode/* - OpenCode CLI
60 # oco/* - OpenCode with Vertex AI
61 # pi/* - Pi coding agent
62 # pir/* - Pi coding agent (read-only/Vertex AI)
63 # review-<tool>/* - Review sessions with specified tool (cr, copilot, opencode, oco, pi, pir)
64 #
65 # Model selection (optional):
66 # claude:opus/home - Use opus model
67 # oco:sonnet/project - Use sonnet model
68 #
69 # Examples:
70 # ssh host/claude/home # Claude in ~/src/home
71 # ssh host/claude:opus/home # Claude with opus in ~/src/home
72 # ssh host/pi/home # Pi coding agent in ~/src/home
73 # ssh host/pi:sonnet/home # Pi with sonnet in ~/src/home
74 # ssh host/pir/home # Pi (Vertex AI) in ~/src/home
75 # ssh host/pir:sonnet/home # Pi (Vertex AI) with sonnet in ~/src/home
76 # ssh host/review-copilot/pr-123="prompt" # Review with copilot
77
78 set -euo pipefail
79
80 # Parse session spec: "claude/home=src/home" -> session="claude/home", dir="src/home"
81 SESSION_SPEC="''${1}"
82 if [[ "$SESSION_SPEC" =~ ^([^=]+)=(.+)$ ]]; then
83 SESSION_NAME="''${BASH_REMATCH[1]}"
84 WORK_DIR="''${BASH_REMATCH[2]}"
85 else
86 SESSION_NAME="$SESSION_SPEC"
87 WORK_DIR=""
88 fi
89
90 # Function to run shpool with a specific command
91 run_with_command() {
92 local session="$1"
93 local cmd="$2"
94 shift 2
95 exec ${pkgs.shpool}/bin/shpool attach -f -c "$cmd $*" "$session"
96 }
97
98 # Function to parse tool:model from prefix
99 # e.g., "claude:opus" -> TOOL=claude, MODEL=opus
100 # "claude" -> TOOL=claude, MODEL=""
101 parse_tool_model() {
102 local prefix="$1"
103 if [[ "$prefix" =~ ^([^:]+):(.+)$ ]]; then
104 TOOL="''${BASH_REMATCH[1]}"
105 MODEL="''${BASH_REMATCH[2]}"
106 else
107 TOOL="$prefix"
108 MODEL=""
109 fi
110 }
111
112 # Function to get default directory from session suffix
113 # e.g., "claude/home" -> "src/home", "claude:opus/home" -> "src/home"
114 get_default_dir() {
115 local session="$1"
116 local prefix="$2"
117 # Remove prefix (which may include :model)
118 local suffix="''${session#$prefix/}"
119 echo "src/$suffix"
120 }
121
122 # Function to run an AI tool session
123 # Args: session_name, tool_prefix, tool_command
124 run_ai_session() {
125 local session="$1"
126 local prefix_pattern="$2"
127 local default_cmd="$3"
128
129 # Extract the prefix part (before first /)
130 local prefix="''${session%%/*}"
131
132 # Parse tool:model from prefix
133 parse_tool_model "$prefix"
134
135 # Use default command if TOOL matches expected pattern
136 local cmd="$default_cmd"
137
138 # Get work directory
139 if [ -z "$WORK_DIR" ]; then
140 WORK_DIR=$(get_default_dir "$session" "$prefix")
141 fi
142
143 run_with_command "$session" "${run-in-dir}" "$cmd" "$WORK_DIR" "$MODEL"
144 }
145
146 # Check session name patterns and run appropriate commands
147 case "$SESSION_NAME" in
148 claude:*/*)
149 # Claude Code with model selection
150 run_ai_session "$SESSION_NAME" "claude" "cr"
151 ;;
152 claude/*)
153 # Claude Code session - run with Vertex AI environment via cr
154 run_ai_session "$SESSION_NAME" "claude" "cr"
155 ;;
156 copilot:*/*)
157 # GitHub Copilot CLI with model selection
158 run_ai_session "$SESSION_NAME" "copilot" "copilot"
159 ;;
160 copilot/*)
161 # GitHub Copilot CLI session
162 run_ai_session "$SESSION_NAME" "copilot" "copilot"
163 ;;
164 opencode:*/*)
165 # OpenCode CLI with model selection
166 run_ai_session "$SESSION_NAME" "opencode" "opencode"
167 ;;
168 opencode/*)
169 # OpenCode CLI session
170 run_ai_session "$SESSION_NAME" "opencode" "opencode"
171 ;;
172 oco:*/*)
173 # OpenCode with Vertex AI and model selection
174 run_ai_session "$SESSION_NAME" "oco" "oco"
175 ;;
176 oco/*)
177 # OpenCode with Vertex AI session
178 run_ai_session "$SESSION_NAME" "oco" "oco"
179 ;;
180 pi:*/*)
181 # Pi coding agent with model selection
182 run_ai_session "$SESSION_NAME" "pi" "pi"
183 ;;
184 pi/*)
185 # Pi coding agent session
186 run_ai_session "$SESSION_NAME" "pi" "pi"
187 ;;
188 pir:*/*)
189 # Pi coding agent (Vertex AI) with model selection
190 run_ai_session "$SESSION_NAME" "pir" "pir"
191 ;;
192 pir/*)
193 # Pi coding agent (Vertex AI) session
194 run_ai_session "$SESSION_NAME" "pir" "pir"
195 ;;
196 review-cr/*|review-claude/*)
197 # Review with Claude Code (cr)
198 if [ -z "$WORK_DIR" ]; then
199 if ${pkgs.shpool}/bin/shpool list 2>/dev/null | grep -q "^$SESSION_NAME "; then
200 exec ${pkgs.shpool}/bin/shpool attach -f "$SESSION_NAME"
201 else
202 echo "Error: new review sessions require a prompt (use session=prompt syntax)" >&2
203 echo "Example: ssh host/review-cr/pr-123=\"Review https://github.com/...\"" >&2
204 exit 1
205 fi
206 fi
207 run_with_command "$SESSION_NAME" "cr" "$WORK_DIR"
208 ;;
209 review-copilot/*)
210 # Review with GitHub Copilot
211 if [ -z "$WORK_DIR" ]; then
212 if ${pkgs.shpool}/bin/shpool list 2>/dev/null | grep -q "^$SESSION_NAME "; then
213 exec ${pkgs.shpool}/bin/shpool attach -f "$SESSION_NAME"
214 else
215 echo "Error: new review sessions require a prompt (use session=prompt syntax)" >&2
216 echo "Example: ssh host/review-copilot/pr-123=\"Review https://github.com/...\"" >&2
217 exit 1
218 fi
219 fi
220 run_with_command "$SESSION_NAME" "copilot" "$WORK_DIR"
221 ;;
222 review-opencode/*)
223 # Review with OpenCode
224 if [ -z "$WORK_DIR" ]; then
225 if ${pkgs.shpool}/bin/shpool list 2>/dev/null | grep -q "^$SESSION_NAME "; then
226 exec ${pkgs.shpool}/bin/shpool attach -f "$SESSION_NAME"
227 else
228 echo "Error: new review sessions require a prompt (use session=prompt syntax)" >&2
229 echo "Example: ssh host/review-opencode/pr-123=\"Review https://github.com/...\"" >&2
230 exit 1
231 fi
232 fi
233 run_with_command "$SESSION_NAME" "opencode" "$WORK_DIR"
234 ;;
235 review-oco/*)
236 # Review with OpenCode Vertex
237 if [ -z "$WORK_DIR" ]; then
238 if ${pkgs.shpool}/bin/shpool list 2>/dev/null | grep -q "^$SESSION_NAME "; then
239 exec ${pkgs.shpool}/bin/shpool attach -f "$SESSION_NAME"
240 else
241 echo "Error: new review sessions require a prompt (use session=prompt syntax)" >&2
242 echo "Example: ssh host/review-oco/pr-123=\"Review https://github.com/...\"" >&2
243 exit 1
244 fi
245 fi
246 run_with_command "$SESSION_NAME" "oco" "$WORK_DIR"
247 ;;
248 review-pi/*)
249 # Review with Pi coding agent
250 if [ -z "$WORK_DIR" ]; then
251 if ${pkgs.shpool}/bin/shpool list 2>/dev/null | grep -q "^$SESSION_NAME "; then
252 exec ${pkgs.shpool}/bin/shpool attach -f "$SESSION_NAME"
253 else
254 echo "Error: new review sessions require a prompt (use session=prompt syntax)" >&2
255 echo "Example: ssh host/review-pi/pr-123=\"Review https://github.com/...\"" >&2
256 exit 1
257 fi
258 fi
259 run_with_command "$SESSION_NAME" "pi" "$WORK_DIR"
260 ;;
261 review-pir/*)
262 # Review with Pi coding agent (Vertex AI)
263 if [ -z "$WORK_DIR" ]; then
264 if ${pkgs.shpool}/bin/shpool list 2>/dev/null | grep -q "^$SESSION_NAME "; then
265 exec ${pkgs.shpool}/bin/shpool attach -f "$SESSION_NAME"
266 else
267 echo "Error: new review sessions require a prompt (use session=prompt syntax)" >&2
268 echo "Example: ssh host/review-pir/pr-123=\"Review https://github.com/...\"" >&2
269 exit 1
270 fi
271 fi
272 run_with_command "$SESSION_NAME" "pir" "$WORK_DIR"
273 ;;
274 *)
275 # If work directory specified, cd into it before attaching
276 if [ -n "$WORK_DIR" ]; then
277 # Start shell and cd into directory
278 exec ${pkgs.shpool}/bin/shpool attach -f -c "${pkgs.zsh}/bin/zsh -c 'cd $WORK_DIR && exec ${pkgs.zsh}/bin/zsh'" "$SESSION_NAME"
279 else
280 # Default: just attach to session with default shell
281 exec ${pkgs.shpool}/bin/shpool attach -f "$SESSION_NAME"
282 fi
283 ;;
284 esac
285''