Update app.py

app.py

@@ -5,12 +5,19 @@ import torch
 import time
 import sys
 import subprocess
-import shlex
+import os
+import pty
+import select
+import signal
+import fcntl
+import struct
+import termios
+import threading
+import re
 from io import StringIO
 import contextlib
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, GenerationConfig
 from typing import List, Tuple
-import os
 from functools import lru_cache
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -18,9 +25,9 @@ os.environ["TOKENIZERS_PARALLELISM"] = "false"
 # --- Configuration ---
 MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 DEVICE = "cpu"
-DTYPE = torch.float32 
+DTYPE = torch.float32
 
-# --- File System (Python Version) ---
+# --- File System ---
 class FileSystem:
     def __init__(self):
         self.files = {
@@ -29,259 +36,273 @@ class FileSystem:
             "notes.txt": "Project requirements:\n- Build a cool AI app\n- Use Gradio"
         }
         self.current_file = "main.py"
-    
+
     def save_file(self, content: str) -> None:
         if self.current_file:
             self.files[self.current_file] = content
-    
+
     def get_current_file_content(self) -> str:
         return self.files.get(self.current_file, "")
-    
+
     def set_current_file(self, filename: str) -> None:
         if filename in self.files:
             self.current_file = filename
-    
+
     def create_file(self, filename: str, content: str = "") -> None:
         if filename not in self.files:
             self.files[filename] = content
             self.current_file = filename
-    
+
     def get_all_files(self) -> List[str]:
         return list(self.files.keys())
 
 fs = FileSystem()
 
-# --- Terminal Engine ---
-class Terminal:
-    """A real interactive terminal that handles pip, shell commands, and python execution."""
-    
-    WELCOME = (
-        "Axon Pro Terminal v1.1\n"
-        "─────────────────────────────────────\n"
-        "  pip install <pkg>    Install packages\n"
-        "  pip uninstall <pkg>  Remove packages\n"
-        "  pip list             Show installed\n"
-        "  python <file.py>     Run a script\n"
-        "  Any shell command    ls, cat, echo, etc.\n"
-        "─────────────────────────────────────\n"
-    )
-    
-    BLOCKED_COMMANDS = [
-        "rm -rf /", "rm -rf /*", "mkfs", "dd if=", ":(){", "fork",
-        "shutdown", "reboot", "halt", "poweroff", "init 0", "init 6",
-    ]
-    
+# --- PTY Terminal Engine ---
+class PTYTerminal:
+    """
+    Real PTY-based terminal. Spawns an actual bash shell so every command
+    that works in a normal Debian/Ubuntu terminal works here — pip, npm,
+    apt-get, git, curl, python, node, you name it.
+    """
+
+    STRIP_ANSI = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07|\x1b\[.*?[@-~]|\r')
+
     def __init__(self):
-        self.history: List[str] = [self.WELCOME]
-        self.env = {}  # Persistent Python globals across exec calls
-    
+        self.master_fd = None
+        self.pid = None
+        self.lock = threading.Lock()
+        self.log_lines: List[str] = []
+        self.max_lines = 800
+        self._spawn_shell()
+        self._append_line("Axon Pro Terminal v2.0 — Full PTY Shell")
+        self._append_line("────────────────────────────────────────")
+        self._append_line("This is a real bash shell. Everything works:")
+        self._append_line("  pip install <pkg>     npm install <pkg>")
+        self._append_line("  apt-get install <pkg> git clone <url>")
+        self._append_line("  python script.py      node app.js")
+        self._append_line("  curl, wget, ls, cat, grep, etc.")
+        self._append_line("────────────────────────────────────────")
+
+    def _spawn_shell(self):
+        """Fork a real bash process attached to a PTY."""
+        try:
+            pid, fd = pty.openpty()
+            self.pid = os.fork()
+            if self.pid == 0:
+                # Child process — become the shell
+                os.close(pid)
+                os.setsid()
+                # Set slave as controlling terminal
+                slave_fd = fd
+                fcntl.ioctl(slave_fd, termios.TIOCSCTTY, 0)
+                os.dup2(slave_fd, 0)
+                os.dup2(slave_fd, 1)
+                os.dup2(slave_fd, 2)
+                if slave_fd > 2:
+                    os.close(slave_fd)
+                env = os.environ.copy()
+                env["TERM"] = "dumb"  # Prevent color codes / fancy prompts
+                env["PS1"] = "$ "     # Simple prompt
+                env["DEBIAN_FRONTEND"] = "noninteractive"
+                os.execvpe("/bin/bash", ["/bin/bash", "--norc", "--noprofile", "-i"], env)
+            else:
+                # Parent process
+                os.close(fd)
+                self.master_fd = pid
+                # Set non-blocking reads
+                flags = fcntl.fcntl(self.master_fd, fcntl.F_GETFL)
+                fcntl.fcntl(self.master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
+                # Set terminal size (cols x rows)
+                winsize = struct.pack("HHHH", 40, 120, 0, 0)
+                fcntl.ioctl(self.master_fd, termios.TIOCSWINSZ, winsize)
+                # Drain initial bash output
+                time.sleep(0.3)
+                self._read_output()
+        except Exception as e:
+            self._append_line(f"[Terminal Error] Failed to spawn shell: {e}")
+            self.master_fd = None
+
+    def _read_output(self, timeout: float = 0.1) -> str:
+        """Read all available output from the PTY."""
+        if self.master_fd is None:
+            return ""
+        output = []
+        deadline = time.time() + timeout
+        while True:
+            remaining = deadline - time.time()
+            if remaining <= 0:
+                break
+            try:
+                ready, _, _ = select.select([self.master_fd], [], [], min(remaining, 0.05))
+                if ready:
+                    chunk = os.read(self.master_fd, 4096)
+                    if chunk:
+                        output.append(chunk.decode("utf-8", errors="replace"))
+                        deadline = time.time() + 0.15  # Extend deadline while data flows
+                    else:
+                        break
+                else:
+                    if output:
+                        break
+            except OSError:
+                break
+        return "".join(output)
+
+    def _clean_output(self, text: str) -> str:
+        """Strip ANSI escape codes and control characters."""
+        cleaned = self.STRIP_ANSI.sub("", text)
+        # Remove null bytes and other control chars (keep newlines/tabs)
+        cleaned = "".join(ch for ch in cleaned if ch == "\n" or ch == "\t" or (ord(ch) >= 32))
+        return cleaned
+
+    def _append_line(self, text: str):
+        self.log_lines.append(text)
+        while len(self.log_lines) > self.max_lines:
+            self.log_lines.pop(0)
+
     def get_log(self) -> str:
-        return "\n".join(self.history)
-    
-    def _append(self, text: str):
-        self.history.append(text)
-        # Keep buffer from growing forever (~500 lines max)
-        total_lines = sum(s.count("\n") + 1 for s in self.history)
-        while total_lines > 500 and len(self.history) > 1:
-            removed = self.history.pop(0)
-            total_lines -= removed.count("\n") + 1
-    
-    def _is_blocked(self, cmd: str) -> bool:
-        cmd_lower = cmd.lower().strip()
-        return any(b in cmd_lower for b in self.BLOCKED_COMMANDS)
-    
+        return "\n".join(self.log_lines)
+
     def run_command(self, cmd: str) -> str:
-        """Process a terminal command and return the full terminal log."""
+        """Send a command to the real bash shell and capture output."""
         cmd = cmd.strip()
         if not cmd:
             return self.get_log()
-        
-        self._append(f"$ {cmd}")
-        
-        if self._is_blocked(cmd):
-            self._append("⛔ Command blocked for safety.")
-            return self.get_log()
-        
-        parts = cmd.split()
-        base = parts[0].lower()
-        
-        if base == "clear":
-            self.history = []
-            return ""
-        elif base == "help":
-            self._append(self.WELCOME)
-        elif base in ("pip", "pip3"):
-            self._handle_pip(parts)
-        elif base in ("python", "python3"):
-            self._handle_python(cmd, parts)
-        else:
-            self._handle_shell(cmd)
-        
-        return self.get_log()
-    
-    def _handle_pip(self, parts: List[str]):
-        """Handle pip install/uninstall/list/show/freeze commands."""
-        if len(parts) < 2:
-            self._append("Usage: pip install <pkg> | pip uninstall <pkg> | pip list | pip show <pkg> | pip freeze")
-            return
-        
-        subcmd = parts[1].lower()
-        
-        if subcmd == "install" and len(parts) >= 3:
-            pkgs = parts[2:]
-            self._append(f"Installing: {' '.join(pkgs)}...")
-            try:
-                result = subprocess.run(
-                    [sys.executable, "-m", "pip", "install"] + pkgs,
-                    capture_output=True, text=True, timeout=180
-                )
-                stdout = result.stdout.strip()
-                stderr = result.stderr.strip()
-                
-                if result.returncode == 0:
-                    for line in stdout.split("\n"):
-                        low = line.lower()
-                        if any(kw in low for kw in ["successfully", "already satisfied", "requirement"]):
-                            self._append(line)
-                    if not any("successfully" in l.lower() or "already satisfied" in l.lower() for l in stdout.split("\n")):
-                        self._append(stdout.split("\n")[-1] if stdout else "Done.")
-                    self._append("✅ pip install completed")
-                else:
-                    err_lines = (stderr or stdout).split("\n")[-6:]
-                    for line in err_lines:
-                        self._append(line)
-                    self._append("❌ pip install failed")
-                    
-            except subprocess.TimeoutExpired:
-                self._append("❌ Timed out (>180s). Try a smaller package or check connection.")
-            except Exception as e:
-                self._append(f"❌ Error: {e}")
-        
-        elif subcmd == "uninstall" and len(parts) >= 3:
-            pkg = parts[2]
-            protected = ["gradio", "torch", "transformers", "pip", "setuptools", "wheel", "numpy"]
-            if pkg.lower() in protected:
-                self._append(f"⛔ '{pkg}' is protected and cannot be uninstalled.")
-                return
-            self._append(f"Uninstalling {pkg}...")
-            try:
-                result = subprocess.run(
-                    [sys.executable, "-m", "pip", "uninstall", "-y", pkg],
-                    capture_output=True, text=True, timeout=60
-                )
-                if result.returncode == 0:
-                    self._append(f"✅ Uninstalled {pkg}")
-                else:
-                    self._append(f"❌ {result.stderr.strip()}")
-            except Exception as e:
-                self._append(f"❌ Error: {e}")
-        
-        elif subcmd == "list":
-            try:
-                result = subprocess.run(
-                    [sys.executable, "-m", "pip", "list", "--format=columns"],
-                    capture_output=True, text=True, timeout=30
-                )
-                self._append(result.stdout.strip() if result.stdout else "No packages found.")
-            except Exception as e:
-                self._append(f"Error: {e}")
-        
-        elif subcmd == "show" and len(parts) >= 3:
+
+        with self.lock:
+            if cmd.lower() == "clear":
+                self.log_lines = []
+                return ""
+
+            if self.master_fd is None:
+                self._append_line(f"$ {cmd}")
+                self._append_line("[Terminal Error] No shell process. Restart the Space.")
+                return self.get_log()
+
+            # Drain any leftover output
+            self._read_output(0.05)
+
+            self._append_line(f"$ {cmd}")
+
+            # Write the command to bash
             try:
-                result = subprocess.run(
-                    [sys.executable, "-m", "pip", "show", parts[2]],
-                    capture_output=True, text=True, timeout=15
-                )
-                if result.returncode == 0:
-                    self._append(result.stdout.strip())
+                os.write(self.master_fd, (cmd + "\n").encode())
+            except OSError as e:
+                self._append_line(f"[Write Error] {e}")
+                return self.get_log()
+
+            # Determine how long to wait based on command type
+            parts = cmd.split()
+            base = parts[0].lower() if parts else ""
+            long_commands = ["pip", "pip3", "npm", "npx", "apt-get", "apt", "git", "wget",
+                             "curl", "make", "cmake", "cargo", "yarn", "conda"]
+
+            if base in long_commands:
+                wait_time = 180  # 3 minutes for installs
+            elif base in ("python", "python3", "node"):
+                wait_time = 60
+            else:
+                wait_time = 15
+
+            # Read output with smart timeout — stop early when command finishes
+            all_output = []
+            start = time.time()
+            idle_count = 0
+            
+            while time.time() - start < wait_time:
+                chunk = self._read_output(0.3)
+                if chunk:
+                    idle_count = 0
+                    all_output.append(chunk)
+                    # Check if we see our prompt back (command finished)
+                    combined = "".join(all_output)
+                    if combined.rstrip().endswith("$"):
+                        break
                 else:
-                    self._append(f"Package '{parts[2]}' not found.")
-            except Exception as e:
-                self._append(f"Error: {e}")
-        
-        elif subcmd == "freeze":
-            try:
-                result = subprocess.run(
-                    [sys.executable, "-m", "pip", "freeze"],
-                    capture_output=True, text=True, timeout=30
-                )
-                self._append(result.stdout.strip() if result.stdout else "No packages.")
-            except Exception as e:
-                self._append(f"Error: {e}")
-        
-        else:
-            self._append(f"pip {subcmd}: unsupported. Use install, uninstall, list, show, freeze.")
-    
-    def _handle_python(self, cmd: str, parts: List[str]):
-        """Handle: python file.py, python -c 'code', python -m module."""
-        if len(parts) < 2:
-            self._append(f"Python {sys.version.split()[0]}")
-            return
+                    idle_count += 1
+                    # For short commands, bail after a brief idle
+                    if base not in long_commands and idle_count >= 3:
+                        break
+                    # For long commands, bail after longer idle
+                    if base in long_commands and idle_count >= 10:
+                        break
+
+            raw = "".join(all_output)
+            cleaned = self._clean_output(raw)
+
+            # Remove the echoed command from output (bash echoes what you type)
+            lines = cleaned.split("\n")
+            filtered = []
+            skip_echo = True
+            for line in lines:
+                stripped = line.strip()
+                if skip_echo and stripped == cmd.strip():
+                    skip_echo = False
+                    continue
+                # Skip bare prompt lines
+                if stripped == "$" or stripped == "$ ":
+                    continue
+                filtered.append(line)
+
+            result = "\n".join(filtered).strip()
+            if result:
+                self._append_line(result)
+
+            return self.get_log()
+
+    def run_editor_code(self, code: str) -> str:
+        """Execute code from the editor by writing it to a temp file and running it."""
+        with self.lock:
+            self._append_line("$ python [editor]")
         
-        if parts[1] == "-c" and len(parts) >= 3:
-            code = cmd.split("-c", 1)[1].strip().strip("'\"")
-            self._exec_python(code)
-        elif parts[1] == "-m" and len(parts) >= 3:
-            try:
-                result = subprocess.run(
-                    [sys.executable] + parts[1:],
-                    capture_output=True, text=True, timeout=30
-                )
-                if result.stdout:
-                    self._append(result.stdout.strip())
-                if result.stderr:
-                    self._append(result.stderr.strip())
-            except subprocess.TimeoutExpired:
-                self._append("Process timed out (30s limit).")
-            except Exception as e:
-                self._append(f"Error: {e}")
-        else:
-            filename = parts[1]
-            if filename in fs.files:
-                self._exec_python(fs.files[filename])
-            elif os.path.exists(filename):
-                with open(filename, "r") as f:
-                    self._exec_python(f.read())
-            else:
-                self._append(f"File not found: {filename}")
-    
-    def _exec_python(self, code: str):
-        """Execute python code with persistent globals and capture output."""
-        buf = StringIO()
-        start = time.time()
+        # Write code to a temp file and run it through the PTY
+        tmp_path = "/tmp/_axon_editor_run.py"
         try:
-            with contextlib.redirect_stdout(buf):
-                exec(code, self.env)
-            elapsed = time.time() - start
-            out = buf.getvalue()
-            if out:
-                self._append(out.rstrip())
-            self._append(f"[Done in {elapsed:.4f}s]")
+            with open(tmp_path, "w") as f:
+                f.write(code)
         except Exception as e:
-            out = buf.getvalue()
-            if out:
-                self._append(out.rstrip())
-            self._append(f"Error: {e}")
-    
-    def _handle_shell(self, cmd: str):
-        """Handle generic shell commands."""
+            with self.lock:
+                self._append_line(f"[Error] Could not write temp file: {e}")
+            return self.get_log()
+
+        # Run through the shell (but bypass the echo/prompt logic — use subprocess for cleaner output)
         try:
             result = subprocess.run(
-                cmd, shell=True,
+                [sys.executable, tmp_path],
                 capture_output=True, text=True, timeout=30
             )
-            if result.stdout:
-                self._append(result.stdout.rstrip())
-            if result.stderr:
-                self._append(result.stderr.rstrip())
-            if result.returncode != 0 and not result.stdout and not result.stderr:
-                self._append(f"Command exited with code {result.returncode}")
+            if result.stdout.strip():
+                with self.lock:
+                    self._append_line(result.stdout.rstrip())
+            if result.stderr.strip():
+                with self.lock:
+                    self._append_line(result.stderr.rstrip())
+            if not result.stdout.strip() and not result.stderr.strip():
+                with self.lock:
+                    self._append_line("(No output)")
         except subprocess.TimeoutExpired:
-            self._append("⏱ Command timed out (30s limit).")
+            with self.lock:
+                self._append_line("[Timed out after 30s]")
         except Exception as e:
-            self._append(f"Error: {e}")
+            with self.lock:
+                self._append_line(f"[Error] {e}")
 
-terminal_engine = Terminal()
+        return self.get_log()
+
+    def cleanup(self):
+        if self.pid and self.pid > 0:
+            try:
+                os.kill(self.pid, signal.SIGTERM)
+            except ProcessLookupError:
+                pass
+        if self.master_fd is not None:
+            try:
+                os.close(self.master_fd)
+            except OSError:
+                pass
+
+terminal = PTYTerminal()
 
 # --- Model Loading ---
 @lru_cache(maxsize=1)
@@ -292,18 +313,21 @@ def load_model():
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
-            torch_dtype=DTYPE,
+            dtype=DTYPE,
             device_map=DEVICE,
             low_cpu_mem_usage=True
         )
+        gen_config = GenerationConfig(
+            max_new_tokens=256,
+            temperature=0.2,
+            top_p=0.95,
+            do_sample=True,
+        )
         pipe = pipeline(
             "text-generation",
             model=model,
             tokenizer=tokenizer,
-            max_new_tokens=256,
-            temperature=0.2,
-            top_p=0.95,
-            do_sample=True
+            generation_config=gen_config,
         )
         print(f"Model loaded in {time.time() - start_time:.2f}s")
         return pipe
@@ -313,61 +337,47 @@ def load_model():
 
 # --- AI Functions ---
 
-def run_editor_code(code: str) -> str:
-    """Run code from the ▶ Run button and pipe output into the terminal."""
-    buf = StringIO()
-    start = time.time()
-    try:
-        with contextlib.redirect_stdout(buf):
-            exec(code, {}, {})
-        elapsed = time.time() - start
-        result = buf.getvalue()
-        if not result:
-            result = "(No output)"
-        terminal_engine._append(f"$ python [editor]\n{result.rstrip()}\n[Done in {elapsed:.4f}s]")
-        return terminal_engine.get_log()
-    except Exception as e:
-        out = buf.getvalue()
-        terminal_engine._append(f"$ python [editor]\n{out}Error: {e}")
-        return terminal_engine.get_log()
-
 def generate_completion(code: str) -> str:
     model = load_model()
-    if not model: return "# Error: Model failed to load."
+    if not model:
+        return "# Error: Model failed to load."
     prompt = f"<|system|>You are a Python expert. Complete the code. No explanations.</s><|user|>\n{code}\n</s><|assistant|>"
     try:
         result = model(prompt, num_return_sequences=1, return_full_text=False)
-        return result[0]['generated_text'].strip()
+        return result[0]["generated_text"].strip()
     except Exception as e:
         return f"# Error: {str(e)}"
 
 def explain_code(code: str) -> str:
     model = load_model()
-    if not model: return "Error loading model."
+    if not model:
+        return "Error loading model."
     prompt = f"<|system|>Explain this Python code concisely.</s><|user|>\n{code}\n</s><|assistant|>"
     try:
         result = model(prompt, max_new_tokens=512)
-        return result[0]['generated_text'].strip()
+        return result[0]["generated_text"].strip()
     except Exception as e:
         return str(e)
 
 def refactor_code(code: str) -> str:
     model = load_model()
-    if not model: return "Error loading model."
+    if not model:
+        return "Error loading model."
     prompt = f"<|system|>Refactor this Python code for PEP 8 and best practices.</s><|user|>\n{code}\n</s><|assistant|>"
     try:
         result = model(prompt, max_new_tokens=512)
-        return result[0]['generated_text'].strip()
+        return result[0]["generated_text"].strip()
     except Exception as e:
         return str(e)
 
 def generate_code(prompt_text: str) -> str:
     model = load_model()
-    if not model: return "Error loading model."
+    if not model:
+        return "Error loading model."
     prompt = f"<|system|>Write Python code for: {prompt_text}</s><|assistant|>"
     try:
         result = model(prompt, max_new_tokens=512)
-        return result[0]['generated_text'].strip()
+        return result[0]["generated_text"].strip()
     except Exception as e:
         return str(e)
 
@@ -379,7 +389,7 @@ def create_diff_view(original: str, modified: str) -> str:
 # --- Gradio UI ---
 
 theme = gr.themes.Default(
-    font=[gr.themes.GoogleFont('JetBrains Mono'), 'monospace'],
+    font=[gr.themes.GoogleFont("JetBrains Mono"), "monospace"],
     primary_hue="blue",
     neutral_hue="gray",
 ).set(
@@ -390,28 +400,33 @@ theme = gr.themes.Default(
 )
 
 css = """
-    * { border-radius: 0 !important; } 
-    .editor-container { height: 50vh; }
-    .terminal-box textarea { 
-        background-color: #0d0d0d !important; 
-        color: #00ff41 !important; 
-        font-family: 'JetBrains Mono', monospace !important; 
+    * { border-radius: 0 !important; }
+    .terminal-box textarea {
+        background-color: #0a0a0a !important;
+        color: #00ff41 !important;
+        font-family: 'JetBrains Mono', monospace !important;
         font-size: 13px !important;
+        line-height: 1.4 !important;
     }
     .terminal-input input {
-        background-color: #0d0d0d !important;
+        background-color: #0a0a0a !important;
         color: #00ff41 !important;
         font-family: 'JetBrains Mono', monospace !important;
-        border: 1px solid #333 !important;
+        border: 1px solid #1a3a1a !important;
+        font-size: 13px !important;
+    }
+    .terminal-input input::placeholder {
+        color: #2a6a2a !important;
     }
     .diff-view { height: 40vh; }
     .code-wrap { white-space: pre-wrap !important; }
+    .status-bar { font-size: 12px; opacity: 0.7; }
 """
 
 with gr.Blocks(title="Axon Pro - Python IDE") as demo:
-    
+
     gr.Markdown("# 🐍 Axon Pro — Python AI IDE")
-    
+
     with gr.Tabs() as main_tabs:
         with gr.Tab("Editor", id="editor-tab"):
             with gr.Row(equal_height=True):
@@ -419,17 +434,19 @@ with gr.Blocks(title="Axon Pro - Python IDE") as demo:
                 with gr.Column(scale=1, min_width=200):
                     gr.Markdown("### 📁 Explorer")
                     file_list = gr.Dropdown(
-                        choices=fs.get_all_files(), 
+                        choices=fs.get_all_files(),
                         value=fs.current_file,
                         label="Files",
                         interactive=True,
-                        container=False
+                        container=False,
                     )
                     with gr.Row():
                         new_file_btn = gr.Button("➕ New", variant="secondary")
                         save_btn = gr.Button("💾 Save", variant="secondary")
-                    new_file_name = gr.Textbox(placeholder="script.py", label="New File", container=False)
-                
+                    new_file_name = gr.Textbox(
+                        placeholder="script.py", label="New File", container=False
+                    )
+
                 # Editor
                 with gr.Column(scale=4):
                     editor = gr.Code(
@@ -438,82 +455,109 @@ with gr.Blocks(title="Axon Pro - Python IDE") as demo:
                         language="python",
                         lines=20,
                         interactive=True,
-                        elem_classes="code-wrap"
+                        elem_classes="code-wrap",
                     )
-                    
+
                     with gr.Row():
                         run_btn = gr.Button("▶ Run", variant="primary")
                         complete_btn = gr.Button("✨ Complete", variant="secondary")
                         explain_btn = gr.Button("📝 Explain", variant="secondary")
                         refactor_btn = gr.Button("🔧 Refactor", variant="secondary")
-                        
+
                     with gr.Tabs():
                         with gr.Tab("Terminal", id="terminal-tab"):
                             terminal_output = gr.Textbox(
-                                value=terminal_engine.get_log(), 
-                                lines=10, 
+                                value=terminal.get_log(),
+                                lines=12,
+                                max_lines=30,
                                 interactive=False,
                                 elem_classes="terminal-box",
                                 label="Terminal",
-                                show_label=False
+                                show_label=False,
                             )
                             with gr.Row():
                                 terminal_input = gr.Textbox(
-                                    placeholder="$ pip install pygame, ls, python main.py ...",
+                                    placeholder="$ type any command: pip install pygame, npm i, git clone, ls ...",
                                     scale=8,
                                     container=False,
-                                    elem_classes="terminal-input"
+                                    elem_classes="terminal-input",
+                                )
+                                terminal_run_btn = gr.Button(
+                                    "⏎", variant="primary", scale=1, min_width=50
                                 )
-                                terminal_run_btn = gr.Button("⏎", variant="primary", scale=1, min_width=50)
-                            clear_btn = gr.Button("Clear", variant="secondary", size="sm")
-                        
+                            clear_btn = gr.Button(
+                                "Clear", variant="secondary", size="sm"
+                            )
+
                         with gr.Tab("AI Chat", id="chat-tab"):
-                            chat_history = gr.Chatbot(label="Axon AI", height=300)
+                            chat_history = gr.Chatbot(
+                                label="Axon AI", height=300, type="messages"
+                            )
                             with gr.Row():
-                                chat_input = gr.Textbox(placeholder="Ask to generate Python code...", scale=7, container=False)
-                                send_btn = gr.Button("Generate", variant="primary", scale=1)
-        
+                                chat_input = gr.Textbox(
+                                    placeholder="Ask to generate Python code...",
+                                    scale=7,
+                                    container=False,
+                                )
+                                send_btn = gr.Button(
+                                    "Generate", variant="primary", scale=1
+                                )
+
         with gr.Tab("Diff View", id="diff-tab"):
-            diff_view = gr.Code(label="AI Changes", language="python", elem_classes="diff-view code-wrap", interactive=False)
+            diff_view = gr.Code(
+                label="AI Changes",
+                language="python",
+                elem_classes="diff-view code-wrap",
+                interactive=False,
+            )
             with gr.Row():
                 apply_btn = gr.Button("Apply Changes", variant="primary")
                 discard_btn = gr.Button("Discard Changes", variant="secondary")
-    
-    status_bar = gr.Markdown("**AXON PRO v1.1** | Python 3.x | CPU Mode | TinyLlama-1.1B | Terminal + pip", elem_classes="status-bar")
-    
+
+    status_bar = gr.Markdown(
+        "**AXON PRO v2.0** | Python 3.x | CPU Mode | TinyLlama-1.1B | Full PTY Terminal",
+        elem_classes="status-bar",
+    )
+
     # State
     current_file_state = gr.State(fs.current_file)
     diff_original_state = gr.State("")
     diff_modified_state = gr.State("")
     diff_mode_state = gr.State(False)
-    
-    # --- Functions ---
+
+    # --- Handler Functions ---
     def update_file(content):
         fs.save_file(content)
         return fs.get_current_file_content()
-    
+
     def load_file(filename):
         fs.set_current_file(filename)
         return fs.get_current_file_content(), filename
-    
+
     def create_file(name):
         if name and "." in name:
             fs.create_file(name)
-            return gr.update(choices=fs.get_all_files(), value=name), fs.get_current_file_content()
+            return (
+                gr.update(choices=fs.get_all_files(), value=name),
+                fs.get_current_file_content(),
+            )
         return gr.update(choices=fs.get_all_files()), fs.get_current_file_content()
-    
+
     def terminal_command(cmd):
-        log = terminal_engine.run_command(cmd)
+        log = terminal.run_command(cmd)
         return log, ""
-    
-    def clear_terminal():
-        terminal_engine.history = []
+
+    def clear_terminal_fn():
+        terminal.log_lines = []
         return ""
-    
+
+    def run_editor_wrapper(content):
+        return terminal.run_editor_code(content)
+
     def complete_wrapper(content):
         comp = generate_completion(content)
         return content + "\n" + comp
-    
+
     def explain_wrapper(content):
         explanation = explain_code(content)
         fs.save_file(content)
@@ -521,61 +565,97 @@ with gr.Blocks(title="Axon Pro - Python IDE") as demo:
         diff_mod = f'"""\nEXPLANATION:\n{explanation}\n"""\n\n{content}'
         diff = create_diff_view(diff_orig, diff_mod)
         return diff, diff_orig, diff_mod, True
-    
+
     def refactor_wrapper(content):
         refactored = refactor_code(content)
         fs.save_file(content)
         diff = create_diff_view(content, refactored)
         return diff, content, refactored, True
 
-    def generate_wrapper(prompt, history):
-        generated = generate_code(prompt)
+    def generate_wrapper(prompt_text, history):
+        generated = generate_code(prompt_text)
         diff = create_diff_view("", generated)
-        new_hist = history + [[prompt, "Generated code available in Diff View"]]
-        return diff, "", generated, True, new_hist, ""
+        new_history = history + [
+            {"role": "user", "content": prompt_text},
+            {"role": "assistant", "content": "Generated code is available in the Diff View tab."},
+        ]
+        return diff, "", generated, True, new_history, ""
 
     # --- Wiring ---
     editor.change(update_file, editor, None)
     file_list.change(load_file, file_list, [editor, current_file_state])
     new_file_btn.click(create_file, new_file_name, [file_list, editor])
-    
-    # Terminal — type commands and hit enter or click ⏎
-    terminal_input.submit(terminal_command, terminal_input, [terminal_output, terminal_input])
-    terminal_run_btn.click(terminal_command, terminal_input, [terminal_output, terminal_input])
-    clear_btn.click(clear_terminal, None, terminal_output)
-    
-    # Run button pipes editor code into terminal
-    run_btn.click(run_editor_code, editor, terminal_output)
+
+    # Terminal
+    terminal_input.submit(
+        terminal_command, terminal_input, [terminal_output, terminal_input]
+    )
+    terminal_run_btn.click(
+        terminal_command, terminal_input, [terminal_output, terminal_input]
+    )
+    clear_btn.click(clear_terminal_fn, None, terminal_output)
+
+    # Run button
+    run_btn.click(run_editor_wrapper, editor, terminal_output)
     complete_btn.click(complete_wrapper, editor, editor)
-    
+
+    # AI features → Diff tab
     explain_btn.click(
-        explain_wrapper, editor, 
-        [diff_view, diff_original_state, diff_modified_state, diff_mode_state]
+        explain_wrapper,
+        editor,
+        [diff_view, diff_original_state, diff_modified_state, diff_mode_state],
     ).then(lambda: gr.Tabs(selected="diff-tab"), None, main_tabs)
-    
+
     refactor_btn.click(
-        refactor_wrapper, editor,
-        [diff_view, diff_original_state, diff_modified_state, diff_mode_state]
+        refactor_wrapper,
+        editor,
+        [diff_view, diff_original_state, diff_modified_state, diff_mode_state],
     ).then(lambda: gr.Tabs(selected="diff-tab"), None, main_tabs)
-    
+
     chat_input.submit(
-        generate_wrapper, [chat_input, chat_history], 
-        [diff_view, diff_original_state, diff_modified_state, diff_mode_state, chat_history, chat_input]
+        generate_wrapper,
+        [chat_input, chat_history],
+        [
+            diff_view,
+            diff_original_state,
+            diff_modified_state,
+            diff_mode_state,
+            chat_history,
+            chat_input,
+        ],
+    ).then(lambda: gr.Tabs(selected="diff-tab"), None, main_tabs)
+
+    send_btn.click(
+        generate_wrapper,
+        [chat_input, chat_history],
+        [
+            diff_view,
+            diff_original_state,
+            diff_modified_state,
+            diff_mode_state,
+            chat_history,
+            chat_input,
+        ],
     ).then(lambda: gr.Tabs(selected="diff-tab"), None, main_tabs)
-    
+
+    # Diff actions
     apply_btn.click(
         lambda mod: (mod, False), diff_modified_state, [editor, diff_mode_state]
     ).then(lambda: gr.Tabs(selected="editor-tab"), None, main_tabs)
-    
+
     discard_btn.click(
         lambda: (gr.update(), False), None, [editor, diff_mode_state]
     ).then(lambda: gr.Tabs(selected="editor-tab"), None, main_tabs)
 
+
 if __name__ == "__main__":
+    import atexit
+    atexit.register(terminal.cleanup)
+
     demo.launch(
-        server_name="0.0.0.0", 
+        server_name="0.0.0.0",
         server_port=7860,
         theme=theme,
         css=css,
-        ssr_mode=False
+        ssr_mode=False,
     )