Update main.py

main.py

@@ -3,6 +3,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 import subprocess
 import asyncio
+import os
 
 from threading import Thread
 from queue import Queue, Empty
@@ -119,3 +120,71 @@ async def websocket_endpoint(websocket: WebSocket, elo: str):
             else:
                 break
         await asyncio.sleep(0.1)
+
+
+
+
+
+########################
+async def get_prompt():
+    # Get current working directory
+    cwd = os.getcwd()
+    # Shorten path if it is too long
+    if len(cwd) > 30:
+        cwd = os.path.basename(cwd)
+        cwd = f".../{cwd}"
+    return f"{cwd}$ "
+
+async def exec_command(command):
+    try:
+        # Execute command using subprocess
+        process = await asyncio.create_subprocess_shell(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            shell=True,
+            executable='/bin/bash'
+        )
+
+        # Wait for command to complete
+        stdout, stderr = await process.communicate()
+
+        # Prepare output to send back
+        output = stdout.decode() + stderr.decode()
+
+    except Exception as e:
+        output = str(e)
+
+    return output
+
+async def send_prompt(websocket):
+    # Send initial prompt to client
+    prompt = await get_prompt()
+    await websocket.send(prompt)
+
+@app.websocket("/shell")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    active_connections.add(websocket)
+    await send_prompt(websocket)
+
+    try:
+        while True:
+            # Receive command from client
+            command = await websocket.receive_text()
+
+            if command == "exit":
+                await websocket.send("Goodbye!")
+                break
+
+            # Execute command
+            output = await exec_command(command)
+
+            # Send output to client
+            await websocket.send(output)
+
+            # Send updated prompt to client
+            await send_prompt(websocket)
+
+    except websockets.exceptions.ConnectionClosedError:
+        pass