import os
import subprocess
import asyncio
from aiohttp import web
# 配置
TTYD_PORT = 7681
SERVER_PORT = 7860
# 简单的 HTML 前端
HTML_CONTENT = """
HF Space Terminal
Initializing...
"""
async def handle_http(request):
return web.Response(text=HTML_CONTENT, content_type='text/html')
async def handle_ws(request):
ws = web.WebSocketResponse(heartbeat=30) # 添加心跳防止超时断开
await ws.prepare(request)
reader = None
writer = None
try:
# 连接到本地运行的 ttyd
reader, writer = await asyncio.open_connection('127.0.0.1', TTYD_PORT)
print("Connected to ttyd backend")
except Exception as e:
print(f"Failed to connect to ttyd: {e}")
await ws.close()
return ws
async def ws_to_tty():
try:
async for msg in ws:
if msg.type == web.WSMsgType.BINARY:
writer.write(msg.data)
elif msg.type == web.WSMsgType.TEXT:
# ttyd 期望字节流，将文本编码为 bytes
writer.write(msg.data.encode('utf-8'))
await writer.drain()
except Exception as e:
print(f"WS->TTY Error: {e}")
finally:
if writer:
writer.close()
try:
await writer.wait_closed()
except:
pass
async def tty_to_ws():
try:
while True:
# 读取 ttyd 的输出
data = await reader.read(4096)
if not data:
break
# 发送给前端 (作为二进制，保留所有控制字符)
await ws.send_bytes(data)
except Exception as e:
print(f"TTY->WS Error: {e}")
finally:
if writer:
writer.close()
try:
await writer.wait_closed()
except:
pass
# 并行运行两个方向的数据流
await asyncio.gather(ws_to_tty(), tty_to_ws())
if not ws.closed:
await ws.close()
return ws
async def on_startup(app):
# 启动 ttyd 子进程
# 使用 --writable 确保可写 (虽然默认通常是可写的，但显式指定更安全)
# 注意：如果之前的 ttyd 版本不支持某些参数，这里只保留最基础的 -p
cmd = ["/usr/local/bin/ttyd", "-p", str(TTYD_PORT), "/bin/bash"]
print(f"Starting ttyd: {' '.join(cmd)}")
# 启动进程
proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
app['ttyd_proc'] = proc
print("ttyd started successfully.")
def main():
app = web.Application()
app.router.add_get('/', handle_http)
app.router.add_get('/ws', handle_ws)
app.on_startup.append(on_startup)
print(f"Starting server on port {SERVER_PORT}...")
# print=None 禁止 aiohttp 自带的日志，避免干扰
web.run_app(app, host='0.0.0.0', port=SERVER_PORT, print=lambda x: None)
if __name__ == '__main__':
main()