|
| 1 | +#!/usr/bin/env python3 |
| 2 | + |
| 3 | +import sys |
| 4 | +import subprocess |
| 5 | +import threading |
| 6 | +import argparse |
| 7 | +import os |
| 8 | + |
| 9 | +# --- Configuration --- |
| 10 | +LOG_FILE = os.path.join(os.path.dirname(os.path.realpath(__file__)), "mcp_io.log") |
| 11 | +# --- End Configuration --- |
| 12 | + |
| 13 | +# --- Argument Parsing --- |
| 14 | +parser = argparse.ArgumentParser( |
| 15 | + description="Wrap a command, passing STDIN/STDOUT verbatim while logging them.", |
| 16 | + usage="%(prog)s <command> [args...]" |
| 17 | +) |
| 18 | +# Capture the command and all subsequent arguments |
| 19 | +parser.add_argument('command', nargs=argparse.REMAINDER, |
| 20 | + help='The command and its arguments to execute.') |
| 21 | + |
| 22 | +open(LOG_FILE, 'w', encoding='utf-8') |
| 23 | + |
| 24 | +if len(sys.argv) == 1: |
| 25 | + parser.print_help(sys.stderr) |
| 26 | + sys.exit(1) |
| 27 | + |
| 28 | +args = parser.parse_args() |
| 29 | + |
| 30 | +if not args.command: |
| 31 | + print("Error: No command provided.", file=sys.stderr) |
| 32 | + parser.print_help(sys.stderr) |
| 33 | + sys.exit(1) |
| 34 | + |
| 35 | +target_command = args.command |
| 36 | +# --- End Argument Parsing --- |
| 37 | + |
| 38 | +# --- I/O Forwarding Functions --- |
| 39 | +# These will run in separate threads |
| 40 | + |
| 41 | +def forward_and_log_stdin(proxy_stdin, target_stdin, log_file): |
| 42 | + """Reads from proxy's stdin, logs it, writes to target's stdin.""" |
| 43 | + try: |
| 44 | + while True: |
| 45 | + # Read line by line from the script's actual stdin |
| 46 | + line_bytes = proxy_stdin.readline() |
| 47 | + if not line_bytes: # EOF reached |
| 48 | + break |
| 49 | + |
| 50 | + # Decode for logging (assuming UTF-8, adjust if needed) |
| 51 | + try: |
| 52 | + line_str = line_bytes.decode('utf-8') |
| 53 | + except UnicodeDecodeError: |
| 54 | + line_str = f"[Non-UTF8 data, {len(line_bytes)} bytes]\n" # Log representation |
| 55 | + |
| 56 | + # Log with prefix |
| 57 | + log_file.write(f"输入: {line_str}") |
| 58 | + log_file.flush() # Ensure log is written promptly |
| 59 | + |
| 60 | + # Write the original bytes to the target process's stdin |
| 61 | + target_stdin.write(line_bytes) |
| 62 | + target_stdin.flush() # Ensure target receives it promptly |
| 63 | + |
| 64 | + except Exception as e: |
| 65 | + # Log errors happening during forwarding |
| 66 | + try: |
| 67 | + log_file.write(f"!!! STDIN Forwarding Error: {e}\n") |
| 68 | + log_file.flush() |
| 69 | + except: pass # Avoid errors trying to log errors if log file is broken |
| 70 | + |
| 71 | + finally: |
| 72 | + # Important: Close the target's stdin when proxy's stdin closes |
| 73 | + # This signals EOF to the target process (like test.sh's read loop) |
| 74 | + try: |
| 75 | + target_stdin.close() |
| 76 | + log_file.write("--- STDIN stream closed to target ---\n") |
| 77 | + log_file.flush() |
| 78 | + except Exception as e: |
| 79 | + try: |
| 80 | + log_file.write(f"!!! Error closing target STDIN: {e}\n") |
| 81 | + log_file.flush() |
| 82 | + except: pass |
| 83 | + |
| 84 | + |
| 85 | +def forward_and_log_stdout(target_stdout, proxy_stdout, log_file): |
| 86 | + """Reads from target's stdout, logs it, writes to proxy's stdout.""" |
| 87 | + try: |
| 88 | + while True: |
| 89 | + # Read line by line from the target process's stdout |
| 90 | + line_bytes = target_stdout.readline() |
| 91 | + if not line_bytes: # EOF reached (process exited or closed stdout) |
| 92 | + break |
| 93 | + |
| 94 | + # Decode for logging |
| 95 | + try: |
| 96 | + line_str = line_bytes.decode('utf-8') |
| 97 | + except UnicodeDecodeError: |
| 98 | + line_str = f"[Non-UTF8 data, {len(line_bytes)} bytes]\n" |
| 99 | + |
| 100 | + # Log with prefix |
| 101 | + log_file.write(f"输出: {line_str}") |
| 102 | + log_file.flush() |
| 103 | + |
| 104 | + # Write the original bytes to the script's actual stdout |
| 105 | + proxy_stdout.write(line_bytes) |
| 106 | + proxy_stdout.flush() # Ensure output is seen promptly |
| 107 | + |
| 108 | + except Exception as e: |
| 109 | + try: |
| 110 | + log_file.write(f"!!! STDOUT Forwarding Error: {e}\n") |
| 111 | + log_file.flush() |
| 112 | + except: pass |
| 113 | + finally: |
| 114 | + try: |
| 115 | + log_file.flush() |
| 116 | + except: pass |
| 117 | + # Don't close proxy_stdout (sys.stdout) here |
| 118 | + |
| 119 | +# --- Main Execution --- |
| 120 | +process = None |
| 121 | +log_f = None |
| 122 | +exit_code = 1 # Default exit code in case of early failure |
| 123 | + |
| 124 | +try: |
| 125 | + # Open log file in append mode ('a') for the threads |
| 126 | + log_f = open(LOG_FILE, 'a', encoding='utf-8') |
| 127 | + |
| 128 | + # Start the target process |
| 129 | + # We use pipes for stdin/stdout |
| 130 | + # We work with bytes (bufsize=0 for unbuffered binary, readline() still works) |
| 131 | + # stderr=subprocess.PIPE could be added to capture stderr too if needed. |
| 132 | + process = subprocess.Popen( |
| 133 | + target_command, |
| 134 | + stdin=subprocess.PIPE, |
| 135 | + stdout=subprocess.PIPE, |
| 136 | + stderr=subprocess.PIPE, # Capture stderr too, good practice |
| 137 | + bufsize=0 # Use 0 for unbuffered binary I/O |
| 138 | + ) |
| 139 | + |
| 140 | + # Pass binary streams to threads |
| 141 | + stdin_thread = threading.Thread( |
| 142 | + target=forward_and_log_stdin, |
| 143 | + args=(sys.stdin.buffer, process.stdin, log_f), |
| 144 | + daemon=True # Allows main thread to exit even if this is stuck (e.g., waiting on stdin) - reconsider if explicit join is needed |
| 145 | + ) |
| 146 | + |
| 147 | + stdout_thread = threading.Thread( |
| 148 | + target=forward_and_log_stdout, |
| 149 | + args=(process.stdout, sys.stdout.buffer, log_f), |
| 150 | + daemon=True |
| 151 | + ) |
| 152 | + |
| 153 | + # Optional: Handle stderr similarly (log and pass through) |
| 154 | + stderr_thread = threading.Thread( |
| 155 | + target=forward_and_log_stdout, # Can reuse the function |
| 156 | + args=(process.stderr, sys.stderr.buffer, log_f), # Pass stderr streams |
| 157 | + # Add a different prefix in the function if needed, or modify function |
| 158 | + # For now, it will log with "STDOUT:" prefix - might want to change function |
| 159 | + # Let's modify the function slightly for this |
| 160 | + daemon=True |
| 161 | + ) |
| 162 | + # A slightly modified version for stderr logging |
| 163 | + def forward_and_log_stderr(target_stderr, proxy_stderr, log_file): |
| 164 | + """Reads from target's stderr, logs it, writes to proxy's stderr.""" |
| 165 | + try: |
| 166 | + while True: |
| 167 | + line_bytes = target_stderr.readline() |
| 168 | + if not line_bytes: break |
| 169 | + try: line_str = line_bytes.decode('utf-8') |
| 170 | + except UnicodeDecodeError: line_str = f"[Non-UTF8 data, {len(line_bytes)} bytes]\n" |
| 171 | + log_file.write(f"STDERR: {line_str}") # Use STDERR prefix |
| 172 | + log_file.flush() |
| 173 | + proxy_stderr.write(line_bytes) |
| 174 | + proxy_stderr.flush() |
| 175 | + except Exception as e: |
| 176 | + try: |
| 177 | + log_file.write(f"!!! STDERR Forwarding Error: {e}\n") |
| 178 | + log_file.flush() |
| 179 | + except: pass |
| 180 | + finally: |
| 181 | + try: |
| 182 | + log_file.flush() |
| 183 | + except: pass |
| 184 | + |
| 185 | + stderr_thread = threading.Thread( |
| 186 | + target=forward_and_log_stderr, |
| 187 | + args=(process.stderr, sys.stderr.buffer, log_f), |
| 188 | + daemon=True |
| 189 | + ) |
| 190 | + |
| 191 | + |
| 192 | + # Start the forwarding threads |
| 193 | + stdin_thread.start() |
| 194 | + stdout_thread.start() |
| 195 | + stderr_thread.start() # Start stderr thread too |
| 196 | + |
| 197 | + # Wait for the target process to complete |
| 198 | + process.wait() |
| 199 | + exit_code = process.returncode |
| 200 | + |
| 201 | + # Wait briefly for I/O threads to finish flushing last messages |
| 202 | + # Since they are daemons, they might exit abruptly with the main thread. |
| 203 | + # Joining them ensures cleaner shutdown and logging. |
| 204 | + # We need to make sure the pipes are closed so the reads terminate. |
| 205 | + # process.wait() ensures target process is dead, pipes should close naturally. |
| 206 | + stdin_thread.join(timeout=1.0) # Add timeout in case thread hangs |
| 207 | + stdout_thread.join(timeout=1.0) |
| 208 | + stderr_thread.join(timeout=1.0) |
| 209 | + |
| 210 | + |
| 211 | +except Exception as e: |
| 212 | + print(f"MCP Logger Error: {e}", file=sys.stderr) |
| 213 | + # Try to log the error too |
| 214 | + if log_f and not log_f.closed: |
| 215 | + try: |
| 216 | + log_f.write(f"!!! MCP Logger Main Error: {e}\n") |
| 217 | + log_f.flush() |
| 218 | + except: pass # Ignore errors during final logging attempt |
| 219 | + exit_code = 1 # Indicate logger failure |
| 220 | + |
| 221 | +finally: |
| 222 | + # Ensure the process is terminated if it's still running (e.g., if logger crashed) |
| 223 | + if process and process.poll() is None: |
| 224 | + try: |
| 225 | + process.terminate() |
| 226 | + process.wait(timeout=1.0) # Give it a moment to terminate |
| 227 | + except: pass # Ignore errors during cleanup |
| 228 | + if process.poll() is None: # Still running? |
| 229 | + try: process.kill() # Force kill |
| 230 | + except: pass # Ignore kill errors |
| 231 | + |
| 232 | + # Final log message |
| 233 | + if log_f and not log_f.closed: |
| 234 | + try: |
| 235 | + log_f.close() |
| 236 | + except: pass # Ignore errors during final logging attempt |
| 237 | + |
| 238 | + # Exit with the target process's exit code |
| 239 | + sys.exit(exit_code) |
0 commit comments