Improve async

This commit is contained in:
luijait 2025-05-25 13:46:30 +00:00
parent d6c620f214
commit fac35329a2
2 changed files with 77 additions and 28 deletions

View File

@ -596,7 +596,7 @@ class OpenAIChatCompletionsModel(Model):
tool_output = None
# For async session inputs with auto_output, suppress the agent message
elif is_async_session_input and has_auto_output:
should_display_message = False
should_display_message = True # Don't suppress async session messages
tool_output = None
# For async session inputs without auto_output, always show the agent message
elif is_async_session_input and not has_auto_output:
@ -3263,15 +3263,41 @@ class _Converter:
# Check if we're in streaming mode
is_streaming_enabled = os.environ.get('CAI_STREAM', 'false').lower() == 'true'
# Always display tool output regardless of streaming mode
cli_print_tool_output(
tool_name=tool_name,
args=tool_args,
output=output_content,
call_id=call_id,
execution_info=execution_info,
token_info=token_info
)
# Check if this output was already displayed during streaming
# For async sessions, we always display since they don't have real streaming
should_display = True
# If streaming is enabled, check if this was already shown
if is_streaming_enabled and hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls:
tool_call_info = cls.recent_tool_calls[call_id]
# Check if this tool was executed very recently (within last 5 seconds)
# This indicates it was likely shown during streaming
if 'start_time' in tool_call_info:
time_since_execution = time.time() - tool_call_info['start_time']
# For regular commands executed recently in streaming mode, skip display
# But always display for async session commands (they have session_id in args)
if time_since_execution < 5.0:
# Parse arguments to check if this is an async session command
try:
import json
args_dict = json.loads(tool_args) if isinstance(tool_args, str) else tool_args
# If it has session_id, it's an async command - always show
if not (isinstance(args_dict, dict) and args_dict.get("session_id")):
should_display = False
except:
# If we can't parse args, assume it's a regular command
should_display = False
# Only display if it hasn't been shown during streaming
if should_display:
cli_print_tool_output(
tool_name=tool_name,
args=tool_args,
output=output_content,
call_id=call_id,
execution_info=execution_info,
token_info=token_info
)
# Continue with normal processing
flush_assistant_message()

View File

@ -174,42 +174,67 @@ class ShellSession: # pylint: disable=too-many-instance-attributes
def _read_output(self):
"""Read output from the process"""
try:
# Buffer for incomplete lines
partial_line = ""
while self.is_running and self.master is not None:
try:
# Check if process has exited before reading
if self.process and self.process.poll() is not None:
self.is_running = False
break
# Read the output
output = os.read(self.master, 1024).decode()
# Read the output - increased buffer size to avoid cutting commands
output = os.read(self.master, 4096).decode('utf-8', errors='replace')
if output:
self.output_buffer.append(output)
# Combine with any partial line from previous read
full_output = partial_line + output
# Split into lines but keep the last partial line if it doesn't end with newline
lines = full_output.split('\n')
# If output doesn't end with newline, the last item is a partial line
if not output.endswith('\n'):
partial_line = lines[-1]
lines = lines[:-1]
else:
partial_line = ""
# Add complete lines to buffer
for line in lines:
if line: # Don't add empty lines
self.output_buffer.append(line)
self.last_activity = time.time()
else:
# os.read() returned empty. This does NOT necessarily mean
# the process itself has exited if self.process.poll() is None.
# It might be idle and waiting for input.
# The self.process.poll() check at the start of the loop iteration
# (or if self.is_process_running() is used in the sleep condition)
# will handle actual termination.
# Thus, we 'pass' here and let the loop continue with its sleep,
# rather than prematurely setting self.is_running = False.
if self.process and self.process.poll() is None:
# Process is alive but PTY read was empty (e.g., idle).
# Do nothing here to change is_running.
pass
else:
# Process is confirmed dead or no process to check,
# and read returned empty. Session is over.
if partial_line:
# Add any remaining partial line
self.output_buffer.append(partial_line)
self.is_running = False
break # Exit the while loop
break
except UnicodeDecodeError as e:
# Handle unicode decode errors gracefully
self.output_buffer.append(f"[Session {self.session_id}] Unicode decode error in output")
continue
except Exception as read_err:
self.output_buffer.append(f"Error reading output buffer: {str(read_err)}")
self.is_running = False
break
self.output_buffer.append(f"Error reading output buffer: {str(read_err)}")
self.is_running = False
break
# Add a small sleep to prevent busy-waiting if no output
if self.is_process_running():
time.sleep(0.05)
time.sleep(0.05)
except Exception as e:
self.output_buffer.append(f"Error in read_output loop: {str(e)}")
self.is_running = False
@ -855,8 +880,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
"command": command,
"args": "",
"session_id": new_session_id,
"async_mode": True,
"creating_session": True # Flag to identify this as session creation
"async_mode": True
}
# Create execution info
@ -1359,8 +1383,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
"command": command,
"args": "",
"session_id": new_session_id,
"async_mode": True,
"creating_session": True # Flag to identify this as session creation
"async_mode": True
}
# Create execution info