1
1
import asyncio
2
2
import os
3
+ import pty
4
+ import sys
3
5
from typing import ClassVar , Literal
4
6
7
+ import pyte
5
8
from anthropic .types .beta import BetaToolBash20241022Param
6
9
7
10
from .base import BaseAnthropicTool , CLIResult , ToolError , ToolResult
@@ -21,20 +24,43 @@ class _BashSession:
21
24
def __init__ (self ):
22
25
self ._started = False
23
26
self ._timed_out = False
27
+ # Create a terminal screen and stream
28
+ self ._screen = pyte .Screen (80 , 24 ) # Standard terminal size
29
+ self ._stream = pyte .Stream (self ._screen )
24
30
25
31
async def start (self ):
26
32
if self ._started :
27
33
return
28
34
29
- self ._process = await asyncio .create_subprocess_shell (
30
- self .command ,
31
- preexec_fn = os .setsid ,
32
- shell = True ,
33
- bufsize = 0 ,
34
- stdin = asyncio .subprocess .PIPE ,
35
- stdout = asyncio .subprocess .PIPE ,
36
- stderr = asyncio .subprocess .PIPE ,
37
- )
35
+ try :
36
+ # Try to create process with PTY
37
+ master , slave = pty .openpty ()
38
+ self ._process = await asyncio .create_subprocess_shell (
39
+ self .command ,
40
+ preexec_fn = os .setsid ,
41
+ shell = True ,
42
+ bufsize = 0 ,
43
+ stdin = asyncio .subprocess .PIPE ,
44
+ stdout = slave ,
45
+ stderr = slave ,
46
+ )
47
+ # Store master fd for reading
48
+ self ._master_fd = master
49
+ self ._using_pty = True
50
+ print ("using pty" )
51
+ except (ImportError , OSError ):
52
+ print ("using pipes" )
53
+ # Fall back to regular pipes if PTY is not available
54
+ self ._process = await asyncio .create_subprocess_shell (
55
+ self .command ,
56
+ preexec_fn = os .setsid ,
57
+ shell = True ,
58
+ bufsize = 0 ,
59
+ stdin = asyncio .subprocess .PIPE ,
60
+ stdout = asyncio .subprocess .PIPE ,
61
+ stderr = asyncio .subprocess .PIPE ,
62
+ )
63
+ self ._using_pty = False
38
64
39
65
self ._started = True
40
66
@@ -45,18 +71,11 @@ def stop(self):
45
71
if self ._process .returncode is not None :
46
72
return
47
73
self ._process .terminate ()
74
+ if hasattr (self , "_master_fd" ):
75
+ os .close (self ._master_fd )
48
76
49
77
async def run (self , command : str ):
50
78
"""Execute a command in the bash shell."""
51
- # Ask for user permission before executing the command
52
- print (f"Do you want to execute the following command?\n { command } " )
53
- user_input = input ("Enter 'yes' to proceed, anything else to cancel: " )
54
-
55
- if user_input .lower () != "yes" :
56
- return ToolResult (
57
- system = "Command execution cancelled by user" ,
58
- error = "User did not provide permission to execute the command." ,
59
- )
60
79
if not self ._started :
61
80
raise ToolError ("Session has not started." )
62
81
if self ._process .returncode is not None :
@@ -71,29 +90,70 @@ async def run(self, command: str):
71
90
72
91
# we know these are not None because we created the process with PIPEs
73
92
assert self ._process .stdin
74
- assert self ._process .stdout
75
- assert self ._process .stderr
76
93
77
94
# send command to the process
78
95
self ._process .stdin .write (
79
96
command .encode () + f"; echo '{ self ._sentinel } '\n " .encode ()
80
97
)
81
98
await self ._process .stdin .drain ()
82
99
83
- # read output from the process, until the sentinel is found
84
100
try :
85
101
async with asyncio .timeout (self ._timeout ):
86
- while True :
87
- await asyncio .sleep (self ._output_delay )
88
- # if we read directly from stdout/stderr, it will wait forever for
89
- # EOF. use the StreamReader buffer directly instead.
90
- output = (
91
- self ._process .stdout ._buffer .decode ()
92
- ) # pyright: ignore[reportAttributeAccessIssue]
93
- if self ._sentinel in output :
94
- # strip the sentinel and break
95
- output = output [: output .index (self ._sentinel )]
96
- break
102
+ if self ._using_pty :
103
+ # Reset screen state
104
+ self ._screen .reset ()
105
+ output = ""
106
+ while True :
107
+ try :
108
+ raw_chunk = os .read (self ._master_fd , 1024 )
109
+ chunk_str = raw_chunk .decode ()
110
+
111
+ # Update output before checking sentinel
112
+ output += chunk_str
113
+
114
+ # Check for sentinel
115
+ if self ._sentinel in chunk_str :
116
+ # Clean the output for display
117
+ clean_chunk = chunk_str [
118
+ : chunk_str .index (self ._sentinel )
119
+ ].encode ()
120
+ if clean_chunk :
121
+ os .write (sys .stdout .fileno (), clean_chunk )
122
+ # Clean the stored output
123
+ if self ._sentinel in output :
124
+ output = output [: output .index (self ._sentinel )]
125
+ break
126
+
127
+ os .write (sys .stdout .fileno (), raw_chunk )
128
+ except OSError :
129
+ break
130
+ await asyncio .sleep (0.01 )
131
+ error = ""
132
+ else :
133
+ # Real-time output for pipe-based reading
134
+ output = ""
135
+ while True :
136
+ chunk = await self ._process .stdout .read (1024 )
137
+ if not chunk :
138
+ break
139
+ chunk_str = chunk .decode ()
140
+ output += chunk_str
141
+
142
+ # Check for sentinel
143
+ if self ._sentinel in chunk_str :
144
+ # Clean the chunk for display
145
+ clean_chunk = chunk_str [
146
+ : chunk_str .index (self ._sentinel )
147
+ ].encode ()
148
+ if clean_chunk :
149
+ os .write (sys .stdout .fileno (), clean_chunk )
150
+ # Clean the stored output
151
+ if self ._sentinel in output :
152
+ output = output [: output .index (self ._sentinel )]
153
+ break
154
+
155
+ os .write (sys .stdout .fileno (), chunk )
156
+ await asyncio .sleep (0.01 )
97
157
except asyncio .TimeoutError :
98
158
self ._timed_out = True
99
159
raise ToolError (
@@ -102,19 +162,24 @@ async def run(self, command: str):
102
162
103
163
if output .endswith ("\n " ):
104
164
output = output [:- 1 ]
105
-
106
- error = (
107
- self ._process .stderr ._buffer .decode ()
108
- ) # pyright: ignore[reportAttributeAccessIssue]
109
- if error .endswith ("\n " ):
165
+ if not self ._using_pty and error .endswith ("\n " ):
110
166
error = error [:- 1 ]
111
167
112
- # clear the buffers so that the next output can be read correctly
113
- self ._process .stdout ._buffer .clear () # pyright: ignore[reportAttributeAccessIssue]
114
- self ._process .stderr ._buffer .clear () # pyright: ignore[reportAttributeAccessIssue]
168
+ # Clear buffers only when using pipes
169
+ if not self ._using_pty :
170
+ self ._process .stdout ._buffer .clear ()
171
+ self ._process .stderr ._buffer .clear ()
115
172
116
173
return CLIResult (output = output , error = error )
117
174
175
+ @staticmethod
176
+ def _strip_ansi (text : str ) -> str :
177
+ """Remove ANSI escape sequences from text."""
178
+ import re
179
+
180
+ ansi_escape = re .compile (r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])" )
181
+ return ansi_escape .sub ("" , text )
182
+
118
183
119
184
class BashTool (BaseAnthropicTool ):
120
185
"""
0 commit comments