""" This module provides a Windows-specific implementation for running commands in a PowerShell session using the pythonnet library to interact with the .NET PowerShell SDK directly. This aims to provide a more robust and integrated way to manage PowerShell processes compared to using temporary script files. """ import os import time import traceback from pathlib import Path from threading import RLock import pythonnet from openhands.core.logger import openhands_logger as logger from openhands.events.action import CmdRunAction from openhands.events.observation import ErrorObservation from openhands.events.observation.commands import ( CmdOutputMetadata, CmdOutputObservation, ) from openhands.runtime.utils.bash_constants import TIMEOUT_MESSAGE_TEMPLATE from openhands.utils.shutdown_listener import should_continue pythonnet.load('coreclr') logger.info("Successfully called pythonnet.load('coreclr')") # Now that pythonnet is initialized, import clr and System try: import clr logger.debug(f'Imported clr module from: {clr.__file__}') # Load System assembly *after* pythonnet is initialized clr.AddReference('System') import System except Exception as clr_sys_ex: raise RuntimeError(f'FATAL: Failed to import clr or System. Error: {clr_sys_ex}') # Attempt to load the PowerShell SDK assembly only if clr and System loaded ps_sdk_path = None try: # Prioritize PowerShell 7+ if available (adjust path if necessary) pwsh7_path = ( Path(os.environ.get('ProgramFiles', 'C:\\Program Files')) / 'PowerShell' / '7' / 'System.Management.Automation.dll' ) if pwsh7_path.exists(): ps_sdk_path = str(pwsh7_path) clr.AddReference(ps_sdk_path) logger.info(f'Loaded PowerShell SDK (Core): {ps_sdk_path}') else: # Fallback to Windows PowerShell 5.1 bundled with Windows winps_path = ( Path(os.environ.get('SystemRoot', 'C:\\Windows')) / 'System32' / 'WindowsPowerShell' / 'v1.0' / 'System.Management.Automation.dll' ) if winps_path.exists(): ps_sdk_path = str(winps_path) clr.AddReference(ps_sdk_path) logger.debug(f'Loaded PowerShell SDK (Desktop): {ps_sdk_path}') else: # Last resort: try loading by assembly name (might work if in GAC or path) clr.AddReference('System.Management.Automation') logger.info( 'Attempted to load PowerShell SDK by name (System.Management.Automation)' ) from System.Management.Automation import JobState, PowerShell from System.Management.Automation.Language import Parser from System.Management.Automation.Runspaces import ( RunspaceFactory, RunspaceState, ) except Exception as e: raise RuntimeError( f'FATAL: Failed to load PowerShell SDK components. Error: {e}. Check pythonnet installation and .NET Runtime compatibility. Path searched: {ps_sdk_path}' ) class WindowsPowershellSession: """ Manages a persistent PowerShell session using the .NET SDK via pythonnet. Allows executing commands within a single runspace, preserving state (variables, current directory) between calls. Handles basic timeout and captures output/error streams. """ def __init__( self, work_dir: str, username: str | None = None, no_change_timeout_seconds: int = 30, max_memory_mb: int | None = None, ): """ Initializes the PowerShell session. Args: work_dir: The starting working directory for the session. username: (Currently ignored) Username for execution. PowerShell SDK typically runs as the current user. no_change_timeout_seconds: Timeout in seconds if no output change is detected (currently NOT fully implemented). max_memory_mb: (Currently ignored) Maximum memory limit for the process. """ # Initialize state flags early to prevent AttributeError in __del__ if init fails self._closed = False self._initialized = False self.runspace = None # Initialize runspace to None if PowerShell is None: # Check if SDK loading failed during module import # Logged critical error during import, just raise here to prevent instantiation raise RuntimeError( 'PowerShell SDK (System.Management.Automation.dll) could not be loaded. Cannot initialize WindowsPowershellSession.' ) self.work_dir = os.path.abspath(work_dir) self.username = username self._cwd = self.work_dir self.NO_CHANGE_TIMEOUT_SECONDS = no_change_timeout_seconds self.max_memory_mb = max_memory_mb # Stored, but not used yet. self.active_job = None self._job_lock = RLock() self._last_job_output = '' # Stores cumulative output returned in the last observation for the active job self._last_job_error: list[ str ] = [] # Stores cumulative errors returned in the last observation for the active job # Create and open the persistent runspace try: # Consider InitialSessionState for more control (e.g., execution policy) # iss = InitialSessionState.CreateDefault() # iss.ExecutionPolicy = Microsoft.PowerShell.ExecutionPolicy.Unrestricted # Requires importing Microsoft.PowerShell namespace # self.runspace = RunspaceFactory.CreateRunspace(iss) self.runspace = RunspaceFactory.CreateRunspace() self.runspace.Open() # Set initial working directory within the runspace self._set_initial_cwd() self._initialized = True # Set to True only on successful initialization logger.info(f'PowerShell runspace created. Initial CWD set to: {self._cwd}') except Exception as e: logger.error(f'Failed to create or open PowerShell runspace: {e}') logger.error(traceback.format_exc()) self.close() # Ensure cleanup if init fails partially raise RuntimeError(f'Failed to initialize PowerShell runspace: {e}') def _set_initial_cwd(self) -> None: """Sets the initial working directory in the runspace.""" ps = None try: ps = PowerShell.Create() ps.Runspace = self.runspace ps.AddScript(f'Set-Location -Path "{self._cwd}"').Invoke() if ps.Streams.Error: errors = '\n'.join([str(err) for err in ps.Streams.Error]) logger.warning(f"Error setting initial CWD to '{self._cwd}': {errors}") # Confirm actual CWD if setting failed self._confirm_cwd() else: logger.debug(f'Successfully set initial runspace CWD to {self._cwd}') # Optional: Confirm CWD even on success for robustness # self._confirm_cwd() except Exception as e: logger.error(f'Exception setting initial CWD: {e}') logger.error(traceback.format_exc()) # Attempt to confirm CWD even if setting threw an exception self._confirm_cwd() finally: if ps: ps.Dispose() def _confirm_cwd(self) -> None: """Confirms the actual CWD in the runspace and updates self._cwd.""" ps_confirm = None try: ps_confirm = PowerShell.Create() ps_confirm.Runspace = self.runspace ps_confirm.AddScript('Get-Location') results = ps_confirm.Invoke() if results and results.Count > 0 and hasattr(results[0], 'Path'): actual_cwd = str(results[0].Path) if os.path.isdir(actual_cwd): if actual_cwd != self._cwd: logger.warning( f'Runspace CWD ({actual_cwd}) differs from expected ({self._cwd}). Updating session CWD.' ) self._cwd = actual_cwd else: logger.debug(f'Confirmed runspace CWD is {self._cwd}') else: logger.error( f'Get-Location returned an invalid path: {actual_cwd}. Session CWD may be inaccurate.' ) elif ps_confirm.Streams.Error: errors = '\n'.join([str(err) for err in ps_confirm.Streams.Error]) logger.error(f'Error confirming runspace CWD: {errors}') else: logger.error('Could not confirm runspace CWD (No result or error).') except Exception as e: logger.error(f'Exception confirming CWD: {e}') finally: if ps_confirm: ps_confirm.Dispose() @property def cwd(self) -> str: """Gets the last known working directory of the session.""" return self._cwd def _run_ps_command( self, script: str, log_output: bool = True ) -> list[System.Management.Automation.PSObject]: """Helper to run a simple synchronous command in the runspace.""" if log_output: logger.debug(f"Running PS command: '{script}'") ps = None results = [] try: ps = PowerShell.Create() ps.Runspace = self.runspace ps.AddScript(script) results = ps.Invoke() except Exception as e: logger.error(f'Exception running script: {script}\n{e}') finally: if ps: ps.Dispose() return results if results else [] def _get_job_object( self, job_id: int | None ) -> System.Management.Automation.Job | None: """Retrieves a job object by its ID.""" script = f'Get-Job -Id {job_id}' results = self._run_ps_command(script, log_output=False) if results and len(results) > 0: potential_job_wrapper = results[0] try: underlying_job = potential_job_wrapper.BaseObject # Basic check for job-like properties before returning _ = underlying_job.Id _ = underlying_job.JobStateInfo.State return underlying_job except AttributeError: logger.warning(f'Retrieved object is not a valid job. ID: {job_id}') return None return None def _receive_job_output( self, job: System.Management.Automation.Job, keep: bool = False ) -> tuple[str, list[str]]: """Receives output and errors from a job.""" if not job: return '', [] output_parts = [] error_parts = [] # Get error stream directly from job object if available try: current_job_obj = self._get_job_object(job.Id) if current_job_obj and current_job_obj.Error: error_records = current_job_obj.Error.ReadAll() if error_records: error_parts.extend([str(e) for e in error_records]) except Exception as read_err: logger.error( f'Failed to read job error stream directly for Job {job.Id}: {read_err}' ) error_parts.append(f'[Direct Error Stream Read Exception: {read_err}]') # Run Receive-Job for the output stream keep_switch = '-Keep' if keep else '' script = f'Receive-Job -Job (Get-Job -Id {job.Id}) {keep_switch}' ps_receive = None try: ps_receive = PowerShell.Create() ps_receive.Runspace = self.runspace ps_receive.AddScript(script) # Collect output results = ps_receive.Invoke() if results: output_parts = [str(r) for r in results] # Collect errors from the Receive-Job command if ps_receive.Streams.Error: receive_job_errors = [str(e) for e in ps_receive.Streams.Error] logger.warning( f'Errors during Receive-Job for Job ID {job.Id}: {receive_job_errors}' ) error_parts.extend(receive_job_errors) except Exception as e: logger.error(f'Exception during Receive-Job for Job ID {job.Id}: {e}') error_parts.append(f'[Receive-Job Exception: {e}]') finally: if ps_receive: ps_receive.Dispose() final_combined_output = '\n'.join(output_parts) return final_combined_output, error_parts def _stop_active_job(self) -> CmdOutputObservation | ErrorObservation: """Stops the active job, collects final output, and cleans up.""" with self._job_lock: job = self.active_job if not job: return ErrorObservation( content='ERROR: No previous running command to interact with.' ) job_id = job.Id # type: ignore[unreachable] logger.info(f'Attempting to stop job ID: {job_id} via C-c.') # Attempt graceful stop stop_script = f'Stop-Job -Job (Get-Job -Id {job_id})' self._run_ps_command(stop_script) # Allow process time to potentially print shutdown messages time.sleep(0.5) # Get final output and errors final_output, final_errors = self._receive_job_output(job, keep=False) combined_output = final_output combined_errors = final_errors # Check job state after stopping final_job = self._get_job_object(job_id) final_state = final_job.JobStateInfo.State if final_job else JobState.Failed logger.info(f'Job {job_id} final state after stop attempt: {final_state}') # Clean up the job remove_script = f'Remove-Job -Job (Get-Job -Id {job_id})' self._run_ps_command(remove_script) # Clear the active job reference self.active_job = None # Construct result output_builder = [combined_output] if combined_output else [] if combined_errors: output_builder.append('\n[ERROR STREAM]') output_builder.extend(combined_errors) # Determine exit code - 0 if Stopped/Completed, 1 otherwise exit_code = ( 0 if final_state in [JobState.Stopped, JobState.Completed] else 1 ) final_content = '\n'.join(output_builder).strip() current_cwd = self._cwd python_safe_cwd = current_cwd.replace('\\\\', '\\\\\\\\') metadata = CmdOutputMetadata( exit_code=exit_code, working_dir=python_safe_cwd ) metadata.suffix = f'\n[The command completed with exit code {exit_code}. CTRL+C was sent.]' return CmdOutputObservation( content=final_content, command='C-c', metadata=metadata, ) def _check_active_job( self, timeout_seconds: int ) -> CmdOutputObservation | ErrorObservation: """ Checks the active job for new output and status, waiting up to timeout_seconds. """ with self._job_lock: if not self.active_job: return ErrorObservation( content='ERROR: No previous running command to retrieve logs from.' ) job_id = self.active_job.Id # type: ignore[unreachable] logger.info( f'Checking active job ID: {job_id} for new output (timeout={timeout_seconds}s).' ) start_time = time.monotonic() monitoring_loop_finished = False accumulated_new_output_builder = [] accumulated_new_errors = [] exit_code = -1 # Assume running final_state = JobState.Running latest_cumulative_output = self._last_job_output latest_cumulative_errors = list(self._last_job_error) while not monitoring_loop_finished: if not should_continue(): logger.warning('Shutdown signal received during job check.') monitoring_loop_finished = True continue elapsed_seconds = time.monotonic() - start_time if elapsed_seconds > timeout_seconds: logger.warning(f'Job check timed out after {timeout_seconds}s.') monitoring_loop_finished = True continue current_job_obj = self._get_job_object(job_id) if not current_job_obj: logger.error(f'Job {job_id} object disappeared during check.') accumulated_new_errors.append('[Job object lost during check]') monitoring_loop_finished = True exit_code = 1 final_state = JobState.Failed if self.active_job and self.active_job.Id == job_id: self.active_job = None continue # Poll output with keep=True (returns cumulative output/errors) polled_cumulative_output, polled_cumulative_errors = ( self._receive_job_output(current_job_obj, keep=True) ) # Detect new output since last poll new_output_detected = '' if polled_cumulative_output != latest_cumulative_output: if polled_cumulative_output.startswith(latest_cumulative_output): new_output_detected = polled_cumulative_output[ len(latest_cumulative_output) : ] else: logger.warning( f'Job {job_id} check: Cumulative output changed unexpectedly' ) new_output_detected = polled_cumulative_output.removeprefix( self._last_job_output ) if new_output_detected.strip(): accumulated_new_output_builder.append( new_output_detected.strip() ) # Detect new errors latest_cumulative_errors_set = set(latest_cumulative_errors) new_errors_detected = [ e for e in polled_cumulative_errors if e not in latest_cumulative_errors_set ] if new_errors_detected: accumulated_new_errors.extend(new_errors_detected) latest_cumulative_output = polled_cumulative_output latest_cumulative_errors = polled_cumulative_errors # Check job state current_state = current_job_obj.JobStateInfo.State if current_state not in [JobState.Running, JobState.NotStarted]: logger.info( f'Job {job_id} finished check loop with state: {current_state}' ) monitoring_loop_finished = True final_state = current_state continue time.sleep(0.1) # Prevent busy-waiting # Process results after loop finished is_finished = final_state not in [JobState.Running, JobState.NotStarted] final_content = '\n'.join(accumulated_new_output_builder).strip() final_errors = list(accumulated_new_errors) if is_finished: logger.info(f'Job {job_id} has finished. Collecting final output.') final_job_obj = self._get_job_object(job_id) if final_job_obj: # Final receive with keep=False to consume remaining output final_cumulative_output, final_cumulative_errors = ( self._receive_job_output(final_job_obj, keep=False) ) # Check for new output in final chunk final_new_output_chunk = '' if final_cumulative_output.startswith(latest_cumulative_output): final_new_output_chunk = final_cumulative_output[ len(latest_cumulative_output) : ] elif final_cumulative_output: final_new_output_chunk = final_cumulative_output.removeprefix( self._last_job_output ) if final_new_output_chunk.strip(): final_content = '\n'.join( filter( None, [final_content, final_new_output_chunk.strip()] ) ) # Check for new errors in final chunk latest_cumulative_errors_set = set(latest_cumulative_errors) new_final_errors = [ e for e in final_cumulative_errors if e not in latest_cumulative_errors_set ] if new_final_errors: final_errors.extend(new_final_errors) # Determine exit code based on state exit_code = 0 if final_state == JobState.Completed else 1 # Clean up job remove_script = f'Remove-Job -Job (Get-Job -Id {job_id})' self._run_ps_command(remove_script) if self.active_job and self.active_job.Id == job_id: self.active_job = None self._last_job_output = '' self._last_job_error = [] else: logger.warning(f'Could not get final job object {job_id}') exit_code = 1 if self.active_job and self.active_job.Id == job_id: self.active_job = None self._last_job_output = '' self._last_job_error = [] else: # Update persistent state with latest cumulative values self._last_job_output = latest_cumulative_output self._last_job_error = list(set(latest_cumulative_errors)) # Append errors to final content if final_errors: error_stream_text = '\n'.join(final_errors) if final_content: final_content += f'\n[ERROR STREAM]\n{error_stream_text}' else: final_content = f'[ERROR STREAM]\n{error_stream_text}' # Ensure exit code is non-zero if errors occurred if exit_code == 0 and final_state != JobState.Completed: exit_code = 1 current_cwd = self._cwd python_safe_cwd = current_cwd.replace('\\\\', '\\\\\\\\') metadata = CmdOutputMetadata( exit_code=exit_code, working_dir=python_safe_cwd ) metadata.prefix = '[Below is the output of the previous command.]\n' if is_finished: metadata.suffix = ( f'\n[The command completed with exit code {exit_code}.]' ) else: metadata.suffix = ( f'\n[The command timed out after {timeout_seconds} seconds. ' f'{TIMEOUT_MESSAGE_TEMPLATE}]' ) return CmdOutputObservation( content=final_content, command='', metadata=metadata, ) def _get_current_cwd(self) -> str: """Gets the current working directory from the runspace.""" # Use helper to run Get-Location results = self._run_ps_command('Get-Location') # --- Add more detailed check logging --- if results and results.Count > 0: # type: ignore[attr-defined] first_result = results[0] has_path_attr = hasattr(first_result, 'Path') if has_path_attr: # Original logic resumes here if hasattr is True fetched_cwd = str(first_result.Path) if os.path.isdir(fetched_cwd): if fetched_cwd != self._cwd: logger.info( f"_get_current_cwd: Fetched CWD '{fetched_cwd}' differs from cached '{self._cwd}'. Updating cache." ) self._cwd = fetched_cwd return self._cwd else: logger.warning( f"_get_current_cwd: Path '{fetched_cwd}' is not a valid directory. Returning cached CWD: {self._cwd}" ) return self._cwd else: # Handle cases where Path attribute is missing (e.g., unexpected object type) # Maybe the path is in BaseObject? try: base_object = first_result.BaseObject if hasattr(base_object, 'Path'): fetched_cwd = str(base_object.Path) if os.path.isdir(fetched_cwd): if fetched_cwd != self._cwd: logger.info( f"_get_current_cwd: Fetched CWD '{fetched_cwd}' (from BaseObject) differs from cached '{self._cwd}'. Updating cache." ) self._cwd = fetched_cwd return self._cwd else: logger.warning( f"_get_current_cwd: Path '{fetched_cwd}' (from BaseObject) is not a valid directory. Returning cached CWD: {self._cwd}" ) return self._cwd else: logger.error( f'_get_current_cwd: BaseObject also lacks Path attribute. Cannot determine CWD from result: {first_result}' ) return self._cwd # Return cached except AttributeError as ae: logger.error( f'_get_current_cwd: Error accessing BaseObject or its Path: {ae}. Result: {first_result}' ) return self._cwd # Return cached except Exception as ex: logger.error( f'_get_current_cwd: Unexpected error checking BaseObject: {ex}. Result: {first_result}' ) return self._cwd # Return cached # This path is taken if _run_ps_command returned [] or results.Count was 0 logger.error( f'_get_current_cwd: No valid results received from Get-Location call. Returning cached CWD: {self._cwd}' ) return self._cwd def execute(self, action: CmdRunAction) -> CmdOutputObservation | ErrorObservation: """ Executes a command, potentially as a PowerShell background job for long-running tasks. Aligned with bash.py behavior regarding command execution and messages. Args: action: The command execution action. Returns: CmdOutputObservation or ErrorObservation. """ if not self._initialized or self._closed: return ErrorObservation( content='PowerShell session is not initialized or has been closed.' ) command = action.command.strip() timeout_seconds = action.timeout or 60 # Default to 60 seconds hard timeout is_input = action.is_input # Check if it's intended as input # Detect if this is a background command (ending with &) run_in_background = False if command.endswith('&'): run_in_background = True command = command[:-1].strip() # Remove the & and extra spaces logger.info(f"Detected background command: '{command}'") logger.info( f"Received command: '{command}', Timeout: {timeout_seconds}s, is_input: {is_input}, background: {run_in_background}" ) # --- Simplified Active Job Handling (aligned with bash.py) --- with self._job_lock: if self.active_job: active_job_obj = self._get_job_object(self.active_job.Id) # type: ignore[unreachable] job_is_finished = False final_output = '' # Initialize before conditional assignment final_errors = [] # Initialize before conditional assignment current_job_state = None # Initialize finished_job_id = ( self.active_job.Id ) # Store ID before potentially clearing self.active_job if active_job_obj: current_job_state = active_job_obj.JobStateInfo.State if current_job_state not in [JobState.Running, JobState.NotStarted]: job_is_finished = True logger.info( f'Active job {finished_job_id} was finished ({current_job_state}) before receiving new command. Cleaning up.' ) # Assign final output/errors here final_output, final_errors = self._receive_job_output( active_job_obj, keep=False ) # Consume final output remove_script = ( f'Remove-Job -Job (Get-Job -Id {finished_job_id})' ) self._run_ps_command(remove_script) # --- Reset persistent state --- self._last_job_output = '' self._last_job_error = [] self.active_job = None # else: job still running, job_is_finished remains False else: # Job object disappeared, consider it finished/gone logger.warning( f'Could not retrieve active job object {finished_job_id}. Assuming finished and clearing.' ) job_is_finished = True current_job_state = ( JobState.Failed ) # Assume failed if object is gone # Assign final output/errors here final_output = '' # No output retrievable final_errors = ['[ERROR: Job object disappeared during check]'] # --- Reset persistent state --- self._last_job_output = '' self._last_job_error = [] self.active_job = None # If the job was found to be finished *during this check*, return its final state now. if job_is_finished: # --- Calculate final new output/errors --- new_output = final_output.removeprefix( self._last_job_output ) # final_output was from keep=False last_error_set = set( self._last_job_error ) # Use the state *before* reset new_errors = [e for e in final_errors if e not in last_error_set] # Construct and return the observation for the completed job using the state captured during cleanup exit_code = 0 if current_job_state == JobState.Completed else 1 output_builder = [new_output] if new_output else [] if new_errors: output_builder.append('\\n[ERROR STREAM]') output_builder.extend(new_errors) content_for_return = '\\n'.join(output_builder).strip() current_cwd = self._cwd # Use cached CWD as job is gone python_safe_cwd = current_cwd.replace('\\\\', '\\\\\\\\') metadata = CmdOutputMetadata( exit_code=exit_code, working_dir=python_safe_cwd ) # Indicate this output is from the *previous* command that just finished. metadata.prefix = ( '[Below is the output of the previous command.]\\n' ) metadata.suffix = ( f'\\n[The command completed with exit code {exit_code}.]' ) logger.info( f"Returning final output for job {finished_job_id} which finished before command '{command}' was processed." ) # Use finished_job_id return CmdOutputObservation( content=content_for_return, command=action.command, # The command that triggered this check (e.g., '') metadata=metadata, ) # If job was NOT finished, check incoming command # This block only runs if the job is still active (job_is_finished is False) if not job_is_finished: if command == '': logger.info( 'Received empty command while job running. Checking job status.' ) # Pass the timeout from the empty command action to _check_active_job return self._check_active_job(timeout_seconds) elif command == 'C-c': logger.info('Received C-c while job running. Stopping job.') return self._stop_active_job() elif is_input: # PowerShell session doesn't directly support stdin injection like bash.py/tmux # This requires a different approach (e.g., named pipes, or specific cmdlets). # For now, return an error indicating this limitation. logger.warning( f"Received input command '{command}' while job active, but direct input injection is not supported in this implementation." ) # Get *new* output since last observation to provide context cumulative_output, cumulative_errors = self._receive_job_output( self.active_job, keep=True ) new_output = cumulative_output.removeprefix( self._last_job_output ) last_error_set = set(self._last_job_error) new_errors = [ e for e in cumulative_errors if e not in last_error_set ] output_builder = [new_output] if new_output else [] if new_errors: output_builder.append('\\n[ERROR STREAM]') output_builder.extend(new_errors) # --- UPDATE persistent state --- # Even though input fails, the user saw this output now self._last_job_output = cumulative_output self._last_job_error = list(set(cumulative_errors)) current_cwd = self._cwd python_safe_cwd = current_cwd.replace('\\\\', '\\\\\\\\') metadata = CmdOutputMetadata( exit_code=-1, working_dir=python_safe_cwd ) # Still running metadata.prefix = ( '[Below is the output of the previous command.]\\n' ) metadata.suffix = ( f"\\n[Your input command '{command}' was NOT processed. Direct input to running processes (is_input=True) " 'is not supported by this PowerShell session implementation. You can use C-c to stop the process.]' ) return CmdOutputObservation( content='\\n'.join(output_builder).strip(), command=action.command, metadata=metadata, ) else: # Any other command arrives while a job is running -> Reject it (bash.py behavior) logger.warning( f"Received new command '{command}' while job {self.active_job.Id} is active. New command NOT executed." ) # Get *new* output since last observation to provide context cumulative_output, cumulative_errors = self._receive_job_output( self.active_job, keep=True ) new_output = cumulative_output.removeprefix( self._last_job_output ) last_error_set = set(self._last_job_error) new_errors = [ e for e in cumulative_errors if e not in last_error_set ] output_builder = [new_output] if new_output else [] if new_errors: output_builder.append('\\n[ERROR STREAM]') output_builder.extend(new_errors) # --- UPDATE persistent state --- # Even though command fails, the user saw this output now self._last_job_output = cumulative_output self._last_job_error = list(set(cumulative_errors)) current_cwd = self._cwd # Use cached CWD python_safe_cwd = current_cwd.replace('\\\\', '\\\\\\\\') metadata = CmdOutputMetadata( exit_code=-1, working_dir=python_safe_cwd ) # Exit code -1 indicates still running metadata.prefix = ( '[Below is the output of the previous command.]\n' ) metadata.suffix = ( f'\n[Your command "{command}" is NOT executed. ' f'The previous command is still running - You CANNOT send new commands until the previous command is completed. ' 'By setting `is_input` to `true`, you can interact with the current process: ' "You may wait longer to see additional output of the previous command by sending empty command '', " 'send other commands to interact with the current process, ' 'or send keys ("C-c", "C-z", "C-d") to interrupt/kill the previous command before sending your new command.]' ) return CmdOutputObservation( content='\\n'.join(output_builder).strip(), command=action.command, # Return the command that was attempted metadata=metadata, ) # --- End Active Job Handling --- # --- If we reach here, there is no active job --- # Handle empty command when NO job is active if command == '': logger.warning('Received empty command string (no active job).') current_cwd = self._get_current_cwd() # Update CWD just in case python_safe_cwd = current_cwd.replace('\\\\', '\\\\\\\\') metadata = CmdOutputMetadata(exit_code=0, working_dir=python_safe_cwd) # Align error message with bash.py error_content = 'ERROR: No previous running command to retrieve logs from.' logger.warning( f'Returning specific error message for empty command: {error_content}' ) # No extra suffix needed # metadata.suffix = f"\n[Empty command received (no active job). CWD: {metadata.working_dir}]" return CmdOutputObservation( content=error_content, command='', metadata=metadata ) # Handle C-* when NO job is active/relevant if command.startswith('C-') and len(command) == 3: logger.warning( f'Received control character command: {command}. Not supported when no job active.' ) current_cwd = self._cwd # Use cached CWD python_safe_cwd = current_cwd.replace('\\\\', '\\\\\\\\') # Align error message with bash.py (no running command to interact with) return ErrorObservation( content='ERROR: No previous running command to interact with.' ) # --- Validate command structure using PowerShell Parser --- # (Keep existing validation logic as it's PowerShell specific and useful) parse_errors = None statements = None try: # Parse the input command string ast, _, parse_errors = Parser.ParseInput(command, None) if parse_errors and parse_errors.Length > 0: error_messages = '\n'.join( [ f' - {err.Message} at Line {err.Extent.StartLineNumber}, Column {err.Extent.StartColumnNumber}' for err in parse_errors ] ) logger.error(f'Command failed PowerShell parsing:\n{error_messages}') return ErrorObservation( content=( f'ERROR: Command could not be parsed by PowerShell.\n' f'Syntax errors detected:\n{error_messages}' ) ) statements = ast.EndBlock.Statements if statements.Count > 1: logger.error( f'Detected {statements.Count} statements in the command. Only one is allowed.' ) # Align error message with bash.py splited_cmds = [ str(s.Extent.Text) for s in statements ] # Try to get text return ErrorObservation( content=( f'ERROR: Cannot execute multiple commands at once.\n' f'Please run each command separately OR chain them into a single command via PowerShell operators (e.g., ; or |).\n' f'Detected commands:\n{"\n".join(f"({i + 1}) {cmd}" for i, cmd in enumerate(splited_cmds))}' ) ) elif statements.Count == 0 and not command.strip().startswith('#'): logger.warning( 'Received command that resulted in zero executable statements (likely whitespace or comment).' ) # Treat as empty command if it parses to nothing return CmdOutputObservation( content='', command=command, metadata=CmdOutputMetadata(exit_code=0, working_dir=self._cwd), ) except Exception as parse_ex: logger.error(f'Exception during PowerShell command parsing: {parse_ex}') logger.error(traceback.format_exc()) return ErrorObservation( content=f'ERROR: An exception occurred while parsing the command: {parse_ex}' ) # --- End validation --- # === Synchronous Execution Path (for CWD commands) === if statements and statements.Count == 1: statement = statements[0] try: from System.Management.Automation.Language import ( CommandAst, PipelineAst, ) # Check PipelineAst if isinstance(statement, PipelineAst): pipeline_elements = statement.PipelineElements if ( pipeline_elements and pipeline_elements.Count == 1 and isinstance(pipeline_elements[0], CommandAst) ): command_ast = pipeline_elements[0] command_name = command_ast.GetCommandName() if command_name and command_name.lower() in [ 'set-location', 'cd', 'push-location', 'pop-location', ]: logger.info( f'execute: Identified CWD command via PipelineAst: {command_name}' ) # Run command and prepare proper CmdOutputObservation ps_results = self._run_ps_command(command) # Get current working directory after CWD command current_cwd = self._get_current_cwd() python_safe_cwd = current_cwd.replace('\\\\', '\\\\\\\\') # Convert results to string output if any output = ( '\n'.join([str(r) for r in ps_results]) if ps_results else '' ) return CmdOutputObservation( content=output, command=command, metadata=CmdOutputMetadata( exit_code=0, working_dir=python_safe_cwd ), ) # Check direct CommandAst elif isinstance(statement, CommandAst): command_name = statement.GetCommandName() if command_name and command_name.lower() in [ 'set-location', 'cd', 'push-location', 'pop-location', ]: logger.info( f'execute: Identified CWD command via direct CommandAst: {command_name}' ) # Run command and prepare proper CmdOutputObservation ps_results = self._run_ps_command(command) # Get current working directory after CWD command current_cwd = self._get_current_cwd() python_safe_cwd = current_cwd.replace('\\\\', '\\\\\\\\') # Convert results to string output if any output = ( '\n'.join([str(r) for r in ps_results]) if ps_results else '' ) return CmdOutputObservation( content=output, command=command, metadata=CmdOutputMetadata( exit_code=0, working_dir=python_safe_cwd ), ) except ImportError as imp_err: logger.error( f'execute: Failed to import CommandAst: {imp_err}. Cannot check for CWD commands.' ) except Exception as ast_err: logger.error(f'execute: Error checking command AST: {ast_err}') # === Asynchronous Execution Path (for non-CWD commands) === logger.info( f"execute: Entering asynchronous execution path for command: '{command}'" ) # --- Start the command as a new asynchronous job --- # Reset state for the new job self._last_job_output = '' self._last_job_error = [] ps_start = None job = None output_builder = [] all_errors = [] exit_code = 1 timed_out = False job_start_failed = False job_id = None try: ps_start = PowerShell.Create() ps_start.Runspace = self.runspace escaped_cwd = self._cwd.replace("'", "''") # Check $? after the command. If it's false, exit 1. start_job_script = f"Start-Job -ScriptBlock {{ Set-Location '{escaped_cwd}'; {command}; if (-not $?) {{ exit 1 }} }}" logger.info(f'Starting command as PowerShell job: {command}') ps_start.AddScript(start_job_script) start_results = ps_start.Invoke() if ps_start.Streams.Error: errors = [str(e) for e in ps_start.Streams.Error] logger.error(f'Errors during Start-Job execution: {errors}') all_errors.extend(errors) ps_get = PowerShell.Create() ps_get.Runspace = self.runspace get_job_script = 'Get-Job | Sort-Object -Property Id -Descending | Select-Object -First 1' ps_get.AddScript(get_job_script) get_results = ps_get.Invoke() if ps_get.Streams.Error: errors = [str(e) for e in ps_get.Streams.Error] logger.error(f'Errors getting latest job: {errors}') all_errors.extend(errors) job_start_failed = True if not job_start_failed and get_results and len(get_results) > 0: potential_job = get_results[0] try: underlying_job = potential_job.BaseObject job_state_test = underlying_job.JobStateInfo.State job = underlying_job job_id = job.Id # For background commands, don't track the job in the session if not run_in_background: with self._job_lock: self.active_job = job logger.info( f'Job retrieved successfully. Job ID: {job.Id}, State: {job_state_test}, Background: {run_in_background}' ) if job_state_test == JobState.Failed: logger.error(f'Job {job.Id} failed immediately after starting.') output_chunk, error_chunk = self._receive_job_output( job, keep=False ) if output_chunk: output_builder.append(output_chunk) if error_chunk: all_errors.extend(error_chunk) job_start_failed = True remove_script = f'Remove-Job -Job (Get-Job -Id {job.Id})' self._run_ps_command(remove_script) with self._job_lock: self.active_job = None except AttributeError as e: logger.error( f'Get-Job returned an object without expected properties on BaseObject: {e}' ) logger.error(traceback.format_exc()) all_errors.append('Get-Job did not return a valid Job object.') job_start_failed = True elif not job_start_failed: logger.error('Get-Job did not return any results.') all_errors.append('Get-Job did not return any results.') job_start_failed = True except Exception as start_ex: logger.error(f'Exception during job start/retrieval: {start_ex}') logger.error(traceback.format_exc()) all_errors.append(f'[Job Start/Get Exception: {start_ex}]') job_start_failed = True finally: if ps_start: ps_start.Dispose() if 'ps_get' in locals() and ps_get: ps_get.Dispose() if job_start_failed: current_cwd = self._get_current_cwd() python_safe_cwd = current_cwd.replace('\\\\', '\\\\\\\\') metadata = CmdOutputMetadata(exit_code=1, working_dir=python_safe_cwd) # Use ErrorObservation for critical failures like job start return ErrorObservation( content='Failed to start PowerShell job.\n[ERRORS]\n' + '\n'.join(all_errors) ) # For background commands, return immediately with success if run_in_background: current_cwd = self._get_current_cwd() python_safe_cwd = current_cwd.replace('\\\\', '\\\\\\\\') metadata = CmdOutputMetadata(exit_code=0, working_dir=python_safe_cwd) metadata.suffix = f'\n[Command started as background job {job_id}.]' return CmdOutputObservation( content=f'[Started background job {job_id}]', command=f'{command} &', metadata=metadata, ) # --- Monitor the Job --- start_time = time.monotonic() monitoring_loop_finished = False shutdown_requested = False final_state = JobState.Failed latest_cumulative_output = ( '' # Tracks the absolute latest cumulative output seen in this loop ) latest_cumulative_errors = [] # Tracks the absolute latest cumulative errors seen in this loop while not monitoring_loop_finished: if not should_continue(): logger.warning('Shutdown signal received during job monitoring.') shutdown_requested = True monitoring_loop_finished = True exit_code = -1 continue elapsed_seconds = time.monotonic() - start_time if elapsed_seconds > timeout_seconds: logger.warning( f'Command job monitoring exceeded timeout ({timeout_seconds}s). Leaving job running.' ) timed_out = True monitoring_loop_finished = True exit_code = -1 continue current_job_obj = self._get_job_object(job_id) if not current_job_obj: logger.error(f'Job {job_id} object disappeared during monitoring.') all_errors.append('[Job object lost during monitoring]') monitoring_loop_finished = True exit_code = 1 final_state = JobState.Failed # Reset state as job is gone self._last_job_output = '' self._last_job_error = [] continue # Poll output (keep=True) -> Returns CUMULATIVE output/errors polled_cumulative_output, polled_cumulative_errors = ( self._receive_job_output(current_job_obj, keep=True) ) # Update the latest cumulative state seen in this loop latest_cumulative_output = polled_cumulative_output latest_cumulative_errors = polled_cumulative_errors # Check job state current_state = current_job_obj.JobStateInfo.State if current_state not in [JobState.Running, JobState.NotStarted]: logger.info( f'Job {job_id} finished monitoring loop with state: {current_state}' ) monitoring_loop_finished = True final_state = current_state continue time.sleep(0.1) # --- Monitoring loop finished --- job_finished_naturally = ( not timed_out and not shutdown_requested and final_state in [JobState.Completed, JobState.Stopped, JobState.Failed] ) determined_cwd = self._cwd final_output_content = '' final_error_content = [] if job_finished_naturally: logger.info( f'Job {job_id} finished naturally with state: {final_state}. Clearing final output buffer.' ) final_cumulative_output = '' final_cumulative_errors: list[str] = [] final_job_obj = self._get_job_object(job_id) if final_job_obj: # Get final output/errors with keep=False final_cumulative_output, final_cumulative_errors = ( self._receive_job_output(final_job_obj, keep=False) ) # Always calculate the output relative to the last observation returned final_output_content = final_cumulative_output.removeprefix( self._last_job_output ) # Also calculate final errors relative to last observation returned last_error_set = set(self._last_job_error) final_error_content = [ e for e in final_cumulative_errors if e not in last_error_set ] else: logger.warning( f'Could not get final job object {job_id} to clear output buffer.' ) # If object is gone, output is what was last seen relative to last observation final_output_content = latest_cumulative_output.removeprefix( self._last_job_output ) last_error_set = set(self._last_job_error) final_error_content = [ e for e in latest_cumulative_errors if e not in last_error_set ] exit_code = 0 if final_state == JobState.Completed else 1 if final_state == JobState.Completed: logger.info(f'Job {job_id} completed successfully. Querying final CWD.') determined_cwd = self._get_current_cwd() else: logger.info( f'Job {job_id} finished but did not complete successfully ({final_state}). Using cached CWD: {self._cwd}' ) determined_cwd = self._cwd with self._job_lock: # Lock to clear active_job remove_script = f'Remove-Job -Job (Get-Job -Id {job_id})' self._run_ps_command(remove_script) self.active_job = None logger.info(f'Cleaned up finished job {job_id}') else: logger.info( f'Job {job_id} did not finish naturally (timeout={timed_out}, shutdown={shutdown_requested}). Using cached CWD: {self._cwd}' ) determined_cwd = self._cwd # Exit code is already -1 from loop exit reason # --- Calculate new output/errors relative to last observation (using latest from loop) --- final_output_content = latest_cumulative_output.removeprefix( self._last_job_output ) final_error_content = [ e for e in latest_cumulative_errors if e not in self._last_job_error ] # --- Update persistent state --- self._last_job_output = latest_cumulative_output self._last_job_error = list( set(latest_cumulative_errors) ) # Store unique errors python_safe_cwd = determined_cwd.replace('\\\\', '\\\\\\\\') # Combine unique output chunks for final observation # Using a set ensures uniqueness if chunks were identical across polls # Join accumulated output_builder parts final_output = final_output_content if final_error_content: # Use the calculated final *new* errors error_stream_text = '\n'.join(final_error_content) if final_output: final_output += f'\n[ERROR STREAM]\n{error_stream_text}' else: final_output = f'[ERROR STREAM]\n{error_stream_text}' if exit_code == 0: # Only check exit code if job finished naturally logger.info( f'Detected errors in stream ({len(final_error_content)} records) but job state was Completed. Forcing exit_code to 1.' ) exit_code = 1 # Create metadata metadata = CmdOutputMetadata(exit_code=exit_code, working_dir=python_safe_cwd) # Determine Suffix if timed_out: # Align suffix with bash.py timeout message suffix = ( f'\n[The command timed out after {timeout_seconds} seconds. ' f'{TIMEOUT_MESSAGE_TEMPLATE}]' ) elif shutdown_requested: # Align suffix with bash.py equivalent (though bash.py might not have specific shutdown message) suffix = f'\n[Command execution cancelled due to shutdown signal. Exit Code: {exit_code}]' elif job_finished_naturally: # Align suffix with bash.py completed message suffix = f'\n[The command completed with exit code {exit_code}.]' else: # Should not happen, but defensive fallback suffix = f'\n[Command execution finished. State: {final_state}, Exit Code: {exit_code}]' metadata.suffix = suffix return CmdOutputObservation( content=final_output, command=command, metadata=metadata ) def close(self) -> None: """Closes the PowerShell runspace and releases resources, stopping any active job.""" if self._closed: return logger.info('Closing PowerShell session runspace.') # Stop and remove any active job before closing runspace with self._job_lock: if self.active_job: logger.warning( # type: ignore[unreachable] f'Session closing with active job {self.active_job.Id}. Attempting to stop and remove.' ) job_id = self.active_job.Id try: # Ensure job object exists before trying to stop/remove active_job_obj = self._get_job_object(job_id) if active_job_obj: stop_script = f'Stop-Job -Job (Get-Job -Id {job_id})' self._run_ps_command( stop_script ) # Use helper before runspace closes time.sleep(0.1) remove_script = f'Remove-Job -Job (Get-Job -Id {job_id})' self._run_ps_command(remove_script) logger.info( f'Stopped and removed active job {job_id} during close.' ) else: logger.warning( f'Could not find job object {job_id} to stop/remove during close.' ) except Exception as e: logger.error( f'Error stopping/removing job {job_id} during close: {e}' ) # --- Reset state even if stop/remove failed --- self._last_job_output = '' self._last_job_error = [] self.active_job = None if hasattr(self, 'runspace') and self.runspace: try: # Check state using System.Management.Automation.Runspaces namespace # Get the state info object first to avoid potential pythonnet issues with nested access runspace_state_info = self.runspace.RunspaceStateInfo if runspace_state_info.State == RunspaceState.Opened: self.runspace.Close() self.runspace.Dispose() logger.info('PowerShell runspace closed and disposed.') except Exception as e: logger.error(f'Error closing/disposing PowerShell runspace: {e}') logger.error(traceback.format_exc()) self.runspace = None self._initialized = False self._closed = True def __del__(self) -> None: """Destructor ensures the runspace is closed.""" self.close()