-
Notifications
You must be signed in to change notification settings - Fork 218
Expand file tree
/
Copy pathtmux_tools.py
More file actions
322 lines (275 loc) · 12.5 KB
/
tmux_tools.py
File metadata and controls
322 lines (275 loc) · 12.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
"""Tmux integration tools for CheetahClaws.
Gives the AI model direct control over tmux sessions: create panes,
send commands, read output, and manage layouts. Auto-detected at
startup — tools are only registered when tmux is available on the host.
"""
import os
import re
import sys
import subprocess
import shlex
import shutil
from tool_registry import ToolDef, register_tool
# ── Detection ────────────────────────────────────────────────────────────────
def _find_tmux() -> str | None:
"""Locate a tmux-compatible binary (tmux on Unix, psmux/tmux.exe on Windows)."""
found = shutil.which("tmux") or shutil.which("psmux")
if found:
return found
if sys.platform == "win32":
candidates = [
os.path.expanduser(r"~\.cargo\bin\psmux.exe"),
os.path.expanduser(r"~\.cargo\bin\tmux.exe"),
]
# Allow override via env var for custom install locations
custom = os.environ.get("CHEETAH_PSMUX_PATH")
if custom:
candidates.insert(0, custom)
# Search common install locations
for base in [os.path.expanduser("~\\Desktop"), os.path.expanduser("~")]:
for name in ("psmux.exe", "tmux.exe"):
p = os.path.join(base, "psmux", "target", "release", name)
candidates.append(p)
p2 = os.path.join(base, "localtest", "psmux", "target", "release", name)
candidates.append(p2)
for c in candidates:
if os.path.isfile(c):
return c
return None
_TMUX_BIN: str | None = _find_tmux()
# Sanitize pattern: only allow alphanumerics, underscores, hyphens, dots, colons
_SAFE_NAME = re.compile(r'^[a-zA-Z0-9_.:-]+$')
# Direction flag constants
_RESIZE_FLAGS = {"up": "-U", "down": "-D", "left": "-L", "right": "-R"}
_READ_ONLY_TOOLS = frozenset(("TmuxListSessions", "TmuxCapture", "TmuxListPanes", "TmuxListWindows"))
def tmux_available() -> bool:
"""Return True if a tmux-compatible binary exists on the system."""
return _TMUX_BIN is not None
def _safe(value: str) -> str:
"""Sanitize a tmux target/session name to prevent shell injection."""
if not value or not _SAFE_NAME.match(value):
raise ValueError(f"Invalid tmux identifier: {value!r}")
return value
def _t(params: dict, key: str = "target") -> str:
"""Build a -t flag from params, or empty string if absent."""
val = params.get(key, "")
return f" -t {_safe(val)}" if val else ""
def _run(cmd: str, timeout: int = 10) -> str:
"""Run a tmux command and return combined stdout+stderr.
Replaces bare 'tmux' prefix with the detected binary path.
Unsets nesting guards ($TMUX / $PSMUX_SESSION) so commands work
from inside an existing session.
"""
try:
if cmd.startswith("tmux "):
cmd = f'"{_TMUX_BIN}" {cmd[5:]}'
env = dict(os.environ)
env.pop("TMUX", None)
env.pop("PSMUX_SESSION", None)
r = subprocess.run(
cmd, shell=True, capture_output=True, text=True,
timeout=timeout, env=env,
)
stdout = r.stdout.strip()
stderr = r.stderr.strip()
if r.returncode != 0 and stderr:
return f"FAILED (exit {r.returncode}): {stderr}"
out = (stdout + ("\n" + stderr if stderr else "")).strip()
return out if out else "(ok)"
except subprocess.TimeoutExpired:
return "Error: tmux command timed out"
except Exception as e:
return f"Error: {e}"
# ── Tool implementations ────────────────────────────────────────────────────
def _tmux_list_sessions(params: dict, config: dict) -> str:
return _run("tmux list-sessions")
def _tmux_new_session(params: dict, config: dict) -> str:
name = _safe(params.get("session_name", "cheetah"))
detach = "-d" if params.get("detached", True) else ""
cmd = params.get("command", "")
shell_part = f" {shlex.quote(cmd)}" if cmd else ""
return _run(f"tmux new-session {detach} -s {name}{shell_part}")
def _tmux_split_window(params: dict, config: dict) -> str:
direction = "-v" if params.get("direction", "vertical") == "vertical" else "-h"
cmd = params.get("command", "")
shell_part = f" {shlex.quote(cmd)}" if cmd else ""
return _run(f"tmux split-window {direction}{_t(params)}{shell_part}")
def _tmux_send_keys(params: dict, config: dict) -> str:
keys = params["keys"]
enter = " Enter" if params.get("press_enter", True) else ""
safe_keys = keys.replace("'", "'\\''")
return _run(f"tmux send-keys{_t(params)} '{safe_keys}'{enter}")
def _tmux_capture_pane(params: dict, config: dict) -> str:
lines = params.get("lines", 50)
return _run(f"tmux capture-pane{_t(params)} -p -S -{int(lines)}")
def _tmux_list_panes(params: dict, config: dict) -> str:
return _run(f"tmux list-panes{_t(params)} -F '#{{pane_index}}: #{{pane_current_command}} [#{{pane_width}}x#{{pane_height}}] #{{?pane_active,(active),}}'")
def _tmux_select_pane(params: dict, config: dict) -> str:
return _run(f"tmux select-pane -t {_safe(params['target'])}")
def _tmux_kill_pane(params: dict, config: dict) -> str:
return _run(f"tmux kill-pane{_t(params)}")
def _tmux_new_window(params: dict, config: dict) -> str:
t_flag = _t(params, "target_session")
name = params.get("window_name", "")
n_flag = f" -n {_safe(name)}" if name else ""
cmd = params.get("command", "")
shell_part = f" {shlex.quote(cmd)}" if cmd else ""
return _run(f"tmux new-window{t_flag}{n_flag}{shell_part}")
def _tmux_list_windows(params: dict, config: dict) -> str:
return _run(f"tmux list-windows{_t(params, 'target_session')} -F '#{{window_index}}: #{{window_name}} [#{{window_width}}x#{{window_height}}] #{{?window_active,(active),}}'")
def _tmux_resize_pane(params: dict, config: dict) -> str:
direction = params.get("direction", "down")
amount = int(params.get("amount", 10))
d_flag = _RESIZE_FLAGS.get(direction, "-D")
return _run(f"tmux resize-pane{_t(params)} {d_flag} {amount}")
# ── Schemas ──────────────────────────────────────────────────────────────────
TMUX_TOOL_SCHEMAS = [
{
"name": "TmuxListSessions",
"description": "List all active tmux sessions.",
"input_schema": {"type": "object", "properties": {}},
},
{
"name": "TmuxNewSession",
"description": "Create a new tmux session. Use detached=true (default) to keep it in the background.",
"input_schema": {
"type": "object",
"properties": {
"session_name": {"type": "string", "description": "Session name (default: cheetah)"},
"detached": {"type": "boolean", "description": "Start detached (default: true)"},
"command": {"type": "string", "description": "Optional command to run in the new session"},
},
},
},
{
"name": "TmuxSplitWindow",
"description": "Split the current tmux pane into two. Creates a new visible terminal pane. (Hint: to run a command and keep the pane open, omit 'command' here and use TmuxSendKeys afterwards).",
"input_schema": {
"type": "object",
"properties": {
"target": {"type": "string", "description": "Target pane (e.g. session:window.pane)"},
"direction": {"type": "string", "enum": ["vertical", "horizontal"], "description": "Split direction (default: vertical)"},
"command": {"type": "string", "description": "Optional command to run in the new pane"},
},
},
},
{
"name": "TmuxSendKeys",
"description": "Send keystrokes/commands to a tmux pane. The command runs visibly in that pane.",
"input_schema": {
"type": "object",
"properties": {
"keys": {"type": "string", "description": "The text or command to send"},
"target": {"type": "string", "description": "Target pane (e.g. session:window.pane)"},
"press_enter": {"type": "boolean", "description": "Press Enter after sending keys (default: true)"},
},
"required": ["keys"],
},
},
{
"name": "TmuxCapture",
"description": "Capture and return the visible text content of a tmux pane. Use this to read command output.",
"input_schema": {
"type": "object",
"properties": {
"target": {"type": "string", "description": "Target pane (e.g. session:window.pane)"},
"lines": {"type": "integer", "description": "Number of history lines to capture (default: 50)"},
},
},
},
{
"name": "TmuxListPanes",
"description": "List all panes in the current session/window with their index, command, and size.",
"input_schema": {
"type": "object",
"properties": {
"target": {"type": "string", "description": "Target session or window"},
},
},
},
{
"name": "TmuxSelectPane",
"description": "Switch focus to a specific tmux pane.",
"input_schema": {
"type": "object",
"properties": {
"target": {"type": "string", "description": "Target pane (e.g. 0, 1, or session:window.pane)"},
},
"required": ["target"],
},
},
{
"name": "TmuxKillPane",
"description": "Close/kill a tmux pane.",
"input_schema": {
"type": "object",
"properties": {
"target": {"type": "string", "description": "Target pane to kill"},
},
},
},
{
"name": "TmuxNewWindow",
"description": "Create a new tmux window (tab) in a session. (Hint: to run a command and keep the window open, omit 'command' here and use TmuxSendKeys afterwards).",
"input_schema": {
"type": "object",
"properties": {
"target_session": {"type": "string", "description": "Session to add the window to"},
"window_name": {"type": "string", "description": "Name for the new window"},
"command": {"type": "string", "description": "Optional command to run"},
},
},
},
{
"name": "TmuxListWindows",
"description": "List all windows in a tmux session.",
"input_schema": {
"type": "object",
"properties": {
"target_session": {"type": "string", "description": "Session name"},
},
},
},
{
"name": "TmuxResizePane",
"description": "Resize a tmux pane in a given direction.",
"input_schema": {
"type": "object",
"properties": {
"target": {"type": "string", "description": "Target pane"},
"direction": {"type": "string", "enum": ["up", "down", "left", "right"], "description": "Resize direction"},
"amount": {"type": "integer", "description": "Number of cells to resize (default: 10)"},
},
},
},
]
# ── Registration ─────────────────────────────────────────────────────────────
_TOOL_FUNCS = {
"TmuxListSessions": _tmux_list_sessions,
"TmuxNewSession": _tmux_new_session,
"TmuxSplitWindow": _tmux_split_window,
"TmuxSendKeys": _tmux_send_keys,
"TmuxCapture": _tmux_capture_pane,
"TmuxListPanes": _tmux_list_panes,
"TmuxSelectPane": _tmux_select_pane,
"TmuxKillPane": _tmux_kill_pane,
"TmuxNewWindow": _tmux_new_window,
"TmuxListWindows": _tmux_list_windows,
"TmuxResizePane": _tmux_resize_pane,
}
def register_tmux_tools() -> int:
"""Register all tmux tools. Returns number of tools registered."""
if not tmux_available():
return 0
schema_map = {s["name"]: s for s in TMUX_TOOL_SCHEMAS}
count = 0
for name, func in _TOOL_FUNCS.items():
register_tool(ToolDef(
name=name,
schema=schema_map[name],
func=func,
read_only=name in _READ_ONLY_TOOLS,
concurrent_safe=True,
))
count += 1
return count