server.py

import http.server
import socketserver
import sys
import os
import subprocess
import json
import llm_call

PORT = 8081

class TraceHandler(http.server.SimpleHTTPRequestHandler):
    def do_POST(self):
        if self.path == '/run':
            content_length = int(self.headers['Content-Length'])
            post_data = self.rfile.read(content_length)
            try:
                # Expect JSON with files dict and entryPoint
                payload = json.loads(post_data.decode('utf-8'))
                files = payload.get('files', {})
                entry_point = payload.get('entryPoint')
                
                if not files or not entry_point:
                    raise ValueError("Missing files or entryPoint")
                
                # Write all files to disk
                file_list = []
                for filename, content in files.items():
                    # Security check: filename should be basename only to prevent path traversal
                    safe_filename = os.path.basename(filename)
                    with open(safe_filename, 'w') as f:
                        f.write(content)
                    file_list.append(safe_filename)
                
                # Run the tracer
                # Make sure entry_point is first
                cmd_args = [sys.executable, 'tracer.py', entry_point]
                for f in file_list:
                    if f != entry_point:
                        cmd_args.append(f)
                
                print(f"Running tracer: {' '.join(cmd_args)}")
                
                result = subprocess.run(
                    cmd_args,
                    capture_output=True,
                    text=True
                )
                
                if result.returncode != 0:
                    print(f"Tracer failed:\n{result.stderr}")
                    self.send_response(500)
                    self.end_headers()
                    self.wfile.write(result.stderr.encode('utf-8'))
                    return
                
                print("Tracer finished successfully.")
                
                # Read the generated JSON data
                try:
                    with open("trace_data.json", "r", encoding='utf-8') as f:
                        trace_json = f.read()
                    
                    self.send_response(200)
                    self.send_header('Content-Type', 'application/json')
                    self.end_headers()
                    self.wfile.write(trace_json.encode('utf-8'))
                except Exception as e:
                    print(f"Error reading trace_data.json: {e}")
                    self.send_response(500)
                    self.end_headers()
                    self.wfile.write(str(e).encode('utf-8'))
                
            except Exception as e:
                print(f"Error: {e}")
                self.send_response(500)
                self.end_headers()
                self.wfile.write(str(e).encode('utf-8'))
        elif self.path == '/explain':
            try:
                content_length = int(self.headers['Content-Length'])
                post_data = self.rfile.read(content_length)
                data = json.loads(post_data.decode('utf-8'))
                
                context_data = {
                    'line_content': data.get('line_content', ''),
                    'line_number': data.get('line_number', 0),
                    'variables': data.get('variables', {}),
                    'stack': data.get('stack', []),
                    'stdout': data.get('stdout', ''),
                    'exec_count': data.get('exec_count', 0),
                    'context_code': data.get('context_code', ''),
                    'error_message': data.get('error_message', None)
                }
                messages = data.get('messages', [])
                
                explanation = llm_call.get_chat_response(messages, context_data)
                
                self.send_response(200)
                self.send_header('Content-type', 'application/json')
                self.end_headers()
                self.wfile.write(json.dumps({'explanation': explanation}).encode('utf-8'))
            except Exception as e:
                print(f"Error in /explain: {e}")
                self.send_response(500)
                self.end_headers()
                self.wfile.write(json.dumps({'error': str(e)}).encode('utf-8'))
        else:
            self.send_error(404)

def main():
    if len(sys.argv) >= 2:
        # Collect all files provided in args
        target_files = sys.argv[1:]
    else:
        # Default to session_script.py
        target_files = ["session_script.py"]
        # Reset session_script.py to empty on startup
        with open("session_script.py", "w") as f:
            f.write("")
    
    # Validate main file (first arg)
    main_file = target_files[0]
    if not os.path.exists(main_file):
        print(f"File not found: {main_file}")
        sys.exit(1)
        
    # Initial trace
    print(f"Generating initial trace for {main_file} (and {len(target_files)-1} other files)...")
    cmd_args = [sys.executable, 'tracer.py'] + target_files
    subprocess.run(cmd_args)
    
    # Set up server
    # Use HTTPServer instead of TCPServer for better HTTP support
    # Allow address reuse to prevent "Address already in use" errors
    http.server.HTTPServer.allow_reuse_address = True
    
    port = PORT
    max_attempts = 10
    httpd = None

    for attempt in range(max_attempts):
        try:
            httpd = http.server.HTTPServer(("", port), TraceHandler)
            break
        except OSError as e:
            if "Address already in use" in str(e) or e.errno == 48:
                print(f"Port {port} is in use, trying {port + 1}...")
                port += 1
            else:
                raise e
    
    if httpd is None:
        print(f"Could not find an open port after {max_attempts} attempts.")
        sys.exit(1)

    with httpd:
        httpd.target_files = target_files # Store list of files
        print(f"Serving at http://localhost:{port}/trace_vis.html")
        try:
            httpd.serve_forever()
        except KeyboardInterrupt:
            print("\nServer stopped.")

if __name__ == "__main__":
    main()