Update app.py

app.py

@@ -1,561 +1,558 @@
-# Generated from trimmed zed.proto
-from google.protobuf import descriptor as _descriptor
-from google.protobuf import descriptor_pool as _descriptor_pool
-from google.protobuf import runtime_version as _runtime_version
-from google.protobuf import symbol_database as _symbol_database
-from google.protobuf.internal import builder as _builder
-_runtime_version.ValidateProtobufRuntimeVersion(
-    _runtime_version.Domain.PUBLIC,
-    5,
-    29,
-    0,
-    '',
-    'zed.proto'
-)
-# @@protoc_insertion_point(imports)
-
-_sym_db = _symbol_database.Default()
-
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\tzed.proto\x12\x0czed.messages\"&\n\x06PeerId\x12\x10\n\x08owner_id\x18\x01 \x01(\r\x12\n\n\x02id\x18\x02 \x01(\r\"\xe2\x05\n\x08\x45nvelope\x12\n\n\x02id\x18\x01 \x01(\r\x12\x1a\n\rresponding_to\x18\x02 \x01(\rH\x01\x88\x01\x01\x12\x35\n\x12original_sender_id\x18\x03 \x01(\x0b\x32\x14.zed.messages.PeerIdH\x02\x88\x01\x01\x12\x14\n\x06\x61\x63k_id\x18\x8a\x02 \x01(\rH\x03\x88\x01\x01\x12$\n\x05hello\x18\x04 \x01(\x0b\x32\x13.zed.messages.HelloH\x00\x12 \n\x03\x61\x63k\x18\x05 \x01(\x0b\x32\x11.zed.messages.AckH\x00\x12$\n\x05\x65rror\x18\x06 \x01(\x0b\x32\x13.zed.messages.ErrorH\x00\x12\"\n\x04ping\x18\x07 \x01(\x0b\x32\x12.zed.messages.PingH\x00\x12\x46\n\x17\x61\x63\x63\x65pt_terms_of_service\x18\xef\x01 \x01(\x0b\x32".zed.messages.AcceptTermsOfServiceH\x00\x12W\n accept_terms_of_service_response\x18\xf0\x01 \x01(\x0b\x32*.zed.messages.AcceptTermsOfServiceResponseH\x00\x12\x33\n\rget_llm_token\x18\xeb\x01 \x01(\x0b\x32\x19.zed.messages.GetLlmTokenH\x00\x12\x44\n\x16get_llm_token_response\x18\xec\x01 \x01(\x0b\x32!.zed.messages.GetLlmTokenResponseH\x00\x42\t\n\x07payloadB\x10\n\x0e_responding_toB\x15\n\x13_original_sender_idB\t\n\x07_ack_idJ\x04\x08W\x10YJ\x06\x08\x9e\x01\x10\xa2\x01J\x06\x08\xa4\x01\x10\xa5\x01J\x06\x08\xa6\x01\x10\xaa\x01J\x06\x08\xb1\x01\x10\xba\x01J\x06\x08\xbc\x01\x10\xbd\x01J\x06\x08\xc1\x01\x10\xc4\x01J\x06\x08\xc5\x01\x10\xc6\x01J\x06\x08\xc8\x01\x10\xcb\x01J\x06\x08\xcd\x01\x10\xcf\x01J\x06\x08\xdd\x01\x10\xde\x01J\x06\x08\xe0\x01\x10\xe6\x01J\x06\x08\xf6\x01\x10\xf7\x01J\x06\x08\xf7\x01\x10\xff\x01J\x06\x08\xff\x01\x10\x81\x02\".\n\x05Hello\x12%\n\x07peer_id\x18\x01 \x01(\x0b\x32\x14.zed.messages.PeerId\"\x06\n\x04Ping\"\x05\n\x03\x41\x63k\"M\n\x05\x45rror\x12\x0f\n\x07message\x18\x01 \x01(\t\x12%\n\x04\x63ode\x18\x02 \x01(\x0e\x32\x17.zed.messages.ErrorCode\x12\x0c\n\x04tags\x18\x03 \x03(\t\"\x16\n\x14\x41\x63\x63\x65ptTermsOfService\"7\n\x1c\x41\x63\x63\x65ptTermsOfServiceResponse\x12\x17\n\x0f\x61\x63\x63\x65pted_tos_at\x18\x01 \x01(\x04\"\r\n\x0bGetLlmToken\"$\n\x13GetLlmTokenResponse\x12\r\n\x05token\x18\x01 \x01(\t*\xef\x02\n\tErrorCode\x12\x0c\n\x08Internal\x10\x00\x12\x11\n\rNoSuchChannel\x10\x01\x12\x10\n\x0c\x44isconnected\x10\x02\x12\r\n\tSignedOut\x10\x03\x12\x13\n\x0fUpgradeRequired\x10\x04\x12\r\n\tForbidden\x10\x05\x12\x0c\n\x08NeedsCla\x10\x07\x12\x13\n\x0fNotARootChannel\x10\x08\x12\x14\n\x10\x42\x61\x64PublicNesting\x10\t\x12\x13\n\x0f\x43ircularNesting\x10\n\x12\x13\n\x0fWrongMoveTarget\x10\x0b\x12\x10\n\x0cUnsharedItem\x10\x0c\x12\x11\n\rNoSuchProject\x10\r\x12$\n DevServerProjectPathDoesNotExist\x10\x10\x12\x19\n\x15RemoteUpgradeRequired\x10\x11\x12\x15\n\x11RateLimitExceeded\x10\x12\x12\x10\n\x0c\x43ommitFailed\x10\x13\"\x04\x08\x06\x10\x06\"\x04\x08\x0e\x10\x0f\x62\x06proto3')
-
-_globals = globals()
-_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
-_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'zed_pb2', _globals)
-if not _descriptor._USE_C_DESCRIPTORS:
-  DESCRIPTOR._loaded_options = None
-  _globals['_ERRORCODE']._serialized_start=1085
-  _globals['_ERRORCODE']._serialized_end=1452
-  _globals['_PEERID']._serialized_start=27
-  _globals['_PEERID']._serialized_end=65
-  _globals['_ENVELOPE']._serialized_start=68
-  _globals['_ENVELOPE']._serialized_end=806
-  _globals['_HELLO']._serialized_start=808
-  _globals['_HELLO']._serialized_end=854
-  _globals['_PING']._serialized_start=856
-  _globals['_PING']._serialized_end=862
-  _globals['_ACK']._serialized_start=864
-  _globals['_ACK']._serialized_end=869
-  _globals['_ERROR']._serialized_start=871
-  _globals['_ERROR']._serialized_end=948
-  _globals['_ACCEPTTERMSOFSERVICE']._serialized_start=950
-  _globals['_ACCEPTTERMSOFSERVICE']._serialized_end=972
-  _globals['_ACCEPTTERMSOFSERVICERESPONSE']._serialized_start=974
-  _globals['_ACCEPTTERMSOFSERVICERESPONSE']._serialized_end=1029
-  _globals['_GETLLMTOKEN']._serialized_start=1031
-  _globals['_GETLLMTOKEN']._serialized_end=1044
-  _globals['_GETLLMTOKENRESPONSE']._serialized_start=1046
-  _globals['_GETLLMTOKENRESPONSE']._serialized_end=1082
-
-# Start of the actual script
-import os
-import json
-import ssl
-import time
-import asyncio
-import logging
-import aiohttp
-from aiohttp import web
-import zstandard as zstd
-from websockets.asyncio.client import connect
-from websockets.exceptions import ConnectionClosed
-import uuid
-
-from google.protobuf.json_format import MessageToDict
-
-Envelope = _sym_db.GetSymbol('zed.messages.Envelope')
-
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(levelname)s: %(message)s'
-)
-logger = logging.getLogger(__name__)
-
-
-
-CONFIG = {
-    "API":{
-        "BASE_URL": "https://zed.dev",
-        "API_KEY": os.getenv("API_KEY","sk-123456"),
-        "BASE_API_URL": "https://collab.zed.dev",
-        "WS_URL": "wss://collab.zed.dev/rpc",
-        "LLM_API_URL": "https://llm.zed.dev/completion",
-    },
-    "LOGIN":{
-        "USER_ID": os.getenv("ZED_USER_ID"),
-        "AUTH": os.getenv("ZED_AUTH_TOKEN")
-    },
-    "SERVER":{
-        "PORT": os.getenv("PORT",5200),
-        "TOKEN_EXPIRY_WARNING_MINUTES": 50  
-    },
-    "MODELS":{
-        "claude-3-5-sonnet-20241022":"claude-3-5-sonnet-latest",
-        "claude-3-7-sonnet-20250219":"claude-3-7-sonnet-20250219"
-    }
-}
-
-highest_message_id = 0
-llm_token = None
-token_timestamp = None
-server_peer_id = None
-active_websocket = None
-proxy_server_running = False
-
-class MessageProcessor:
-    @staticmethod
-    def create_chat_response(message, model, is_stream=False):
-        base_response = {
-            "id": f"chatcmpl-{uuid.uuid4()}",
-            "created": int(time.time()),
-            "model": model
-        }
-
-        if is_stream:
-            return {
-                **base_response,
-                "object": "chat.completion.chunk",
-                "choices": [{
-                    "index": 0,
-                    "delta": {
-                        "content": message
-                    }
-                }]
-            }
-
-        return {
-            **base_response,
-            "object": "chat.completion",
-            "choices": [{
-                "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": message
-                },
-                "finish_reason": "stop"
-            }],
-            "usage": None
-        }
-
-def decode_envelope(data):
-    try:
-        dctx = zstd.ZstdDecompressor()
-        decompressed_data = b''
-        with dctx.stream_reader(data) as reader:
-            while True:
-                chunk = reader.read(8192)
-                if not chunk:
-                    break
-                decompressed_data += chunk
-        envelope = Envelope()
-        envelope.ParseFromString(decompressed_data)
-        return MessageToDict(envelope, preserving_proto_field_name=True)
-    except Exception as e:
-        hex_preview = ' '.join(f'{byte:02x}' for byte in data[:20]) + ('...' if len(data) > 20 else '')
-        logger.error(f"无法解码消息： {e}; 数据预览： {hex_preview}")
-        return {"error": f"无法解码消息： {e}"}
-
-def compress_protobuf(data):
-    return zstd.ZstdCompressor(level=-7).compress(data)
-
-def create_message(message_type):
-    global highest_message_id
-    highest_message_id += 1
-    message_id = highest_message_id
-    envelope = Envelope(id=highest_message_id)
-
-    getattr(envelope, message_type).SetInParent()
-    return compress_protobuf(envelope.SerializeToString()), message_id
-
-async def ping_periodically(websocket):
-    while True:
-        try:
-            await websocket.ping()
-            await asyncio.sleep(1)
-        except Exception as e:
-            logger.error(f"发送ping错误: {e}")
-            break
-
-async def handle_messages(websocket):
-    global server_peer_id, llm_token, token_timestamp, active_websocket
-    active_websocket = websocket
-    try:
-        async for message in websocket:
-            message_bytes = message.encode('utf-8') if isinstance(message, str) else message
-            decoded = decode_envelope(message_bytes)
-            if "hello" in decoded:
-                server_peer_id = decoded.get('hello', {}).get('peer_id')
-            elif "accept_terms_of_service_response" in decoded:
-                await request_llm_token(websocket)
-            elif ("get_llm_token_response" in decoded and 
-                  'token' in decoded.get('get_llm_token_response', {})):
-                llm_token = decoded['get_llm_token_response']['token']
-                token_timestamp = time.time()
-                logger.info(f"LLM令牌收到  {time.ctime(token_timestamp)}")
-                if not proxy_server_running:
-                    asyncio.create_task(start_proxy_server())
-                asyncio.create_task(monitor_token_expiration())
-                logger.info("关闭WebSocket连接，直到需要刷新令牌")
-                await websocket.close()
-                active_websocket = None
-                return
-    except ConnectionClosed:
-        logger.info("连接已关闭")
-        active_websocket = None
-
-async def request_llm_token(websocket):
-    message, _ = create_message('get_llm_token')
-    logger.info("请求LLM令牌")
-    await websocket.send(message)
-
-async def request_accept_terms_of_service(websocket):
-    message, _ = create_message('accept_terms_of_service')
-    logger.info("发送同意Zed服务条款")
-    await websocket.send(message)
-
-def format_content(content):
-    if isinstance(content, str):
-        return [{"type": "text", "text": content}]
-    return content
-
-
-
-async def process_message_content(content):
-  """
-  处理消息内容，将不同类型的内容转换为字符串
-  """
-  if isinstance(content, str):
-      return content
-  
-  if isinstance(content, list):
-      return '\n'.join([item.get('text', '') for item in content])
-  
-  if isinstance(content, dict):
-      return content.get('text', None)
-  
-  return None
-
-async def transform_messages(request):
-  """
-  转换消息格式，合并系统消息并处理消息结构
-  """
-  system_message = ''  # 存储系统消息的变量
-  is_collecting_system_message = False  # 是否正在收集系统消息
-  has_processed_system_messages = False  # 是否已处理初始系统消息
-
-  converted_messages = []
-
-  for current in request.get('messages', []):
-      role = current.get('role')
-      current_content = await process_message_content(current.get('content'))
-
-      if current_content is None:
-          converted_messages.append(current)
-          continue
-
-      if role == 'system' and not has_processed_system_messages:
-          if not is_collecting_system_message:
-              # 第一次遇到system，开启收集
-              system_message = current_content
-              is_collecting_system_message = True
-          else:
-              # 继续遇到system，合并system消息
-              system_message += '\n' + current_content
-          continue
-
-      # 遇到非system消息
-      if is_collecting_system_message:
-          # 结束系统消息收集
-          is_collecting_system_message = False
-          has_processed_system_messages = True
-
-      # 如果已处理初始消息序列且再次遇到system，则转换role为user
-      if has_processed_system_messages and role == 'system':
-          role = 'user'
-
-      # 检查是否可以合并消息
-      if converted_messages and converted_messages[-1].get('role') == role:
-          converted_messages[-1]['content'][0]['text'] += '\r\n' + current_content
-      else:
-          converted_messages.append({
-              'role': role,
-              'content': [{'type': 'text', 'text': current_content}]
-          })
-
-  return {
-      'messages': converted_messages,
-      'system': system_message,
-      'model': CONFIG['MODELS'].get(request.get('model'), "claude-3-5-sonnet-latest"),
-      'max_tokens': request.get('max_tokens',8192),
-      'temperature': max(0, min(request.get('temperature', 0), 1)),
-      'top_p': max(0, min(request.get('top_p', 1), 1)),
-      'top_k': max(0, min(request.get('top_k', 0), 500)),
-      'stream': True
-  }
-
-@web.middleware
-async def auth_middleware(request, handler):
-    if CONFIG['API']['API_KEY']:
-        auth_header = request.headers.get('Authorization')
-        xapi_key_header = request.headers.get('x-api-key')
-
-        auth_password = None
-        if auth_header and auth_header.startswith('Bearer '):
-            auth_password = auth_header[7:]
-
-        if auth_password == CONFIG['API']['API_KEY'] or xapi_key_header == CONFIG['API']['API_KEY']:
-            return await handler(request)
-        else:
-            return web.json_response(
-                {"error": "Unauthorized"},
-                status=401
-            )
-
-    return await handler(request)
-
-async def handle_models_request(request):
-    return web.json_response({
-        "object": "list",
-        "data": [
-            {
-                "id": model,
-                "object": "model",
-                "created": int(time.time()),
-                "owned_by": "zed"
-            }
-            for model in CONFIG["MODELS"].keys()
-        ]
-    })
-
-async def handle_message_request(request):
-    global llm_token
-    if not llm_token:
-        return web.json_response({"error": "LLM令牌不可用"}, status=500)
-    try:
-        body = await request.json()
-        isClaudeAI = False
-        if request.path == '/v1/messages':
-            isClaudeAI = True
-            if "messages" in body:
-                for msg in body["messages"]:
-                    if "content" in msg:
-                        msg["content"] = format_content(msg["content"])
-            if "system" in body:
-                if isinstance(body["system"], list):
-                    body["system"] = "\n".join([item["text"] for item in body["system"]])
-            if "model" in body:
-                body["model"] = CONFIG['MODELS'].get(body["model"], "claude-3-5-sonnet-latest")
-        else:
-            body = await transform_messages(body)
-        with open('request_payload222.json', 'w', encoding='utf-8') as f:
-            json.dump(body, f, ensure_ascii=False, indent=2)
-        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {llm_token}"}
-        with open('ceshi.txt', 'w', encoding='utf-8') as f:
-            f.write(llm_token + '\n')
-        payload = {
-            "provider": "anthropic",
-            "model": body.get("model", "claude-3-5-sonnet-latest"),
-            "provider_request": body
-        }
-        # with open('ceshi.txt', 'w', encoding='utf-8') as f:
-        #     f.write(json.dumps(body,ensure_ascii=False) + '\n')
-        if body.get("stream", False):
-            return await handle_streaming_request(request, headers, payload, isClaudeAI)
-        else:
-            return await handle_non_streaming_request(headers, payload, isClaudeAI)
-    except Exception as e:
-        logger.error(f"处理请求时发生错误: {e}")
-        return web.json_response({"error": str(e)}, status=500)
-
-async def handle_non_streaming_request(headers, payload, isClaudeAI=False):
-    async with aiohttp.ClientSession() as session:
-        async with session.post(CONFIG['API']['LLM_API_URL'], headers=headers, json=payload) as r:
-            if r.status != 200:
-                text = await r.text()
-                logger.error(f"LLM API错误: {text}")
-                return web.json_response({"error": text}, status=r.status)
-            full_content, message_data = "", {}
-            async for line in r.content:
-                if not line:
-                    continue
-                try:
-                    event = json.loads(line.decode('utf-8').strip())
-                    et = event.get('type')
-                    if et == "message_start":
-                        message_data = event.get('message', {})
-                    elif et == "content_block_delta" and event.get('delta', {}).get('type') == "text_delta":
-                        full_content += event['delta'].get('text', '')
-                    elif et == "message_delta" and 'usage' in event:
-                        message_data['usage'] = event.get('usage')
-                    elif et == "message_stop":
-                        break
-                except Exception as e:
-                    logger.error(f"Error processing line: {e}")
-            if isClaudeAI:
-                message_data['content'] = [{"type": "text", "text": full_content}]
-            else:
-                message_data = MessageProcessor.create_chat_response(full_content, payload.get("model"), False)
-            return web.json_response(message_data)
-
-async def handle_streaming_request(request, headers, payload, isClaudeAI=False):
-    response = web.StreamResponse()
-    response.headers['Content-Type'] = 'text/event-stream'
-    response.headers['Cache-Control'] = 'no-cache'
-    response.headers['Connection'] = 'keep-alive'
-    await response.prepare(request)
-    logger.info(f"开始处理流请求")
-    async with aiohttp.ClientSession() as session:
-        async with session.post(CONFIG['API']['LLM_API_URL'], headers=headers, json=payload) as api_response:
-            if api_response.status != 200:
-                error_text = await api_response.text()
-                logger.error(f"LLM API (stream)错误: {error_text}")
-                await response.write(f"data: {json.dumps({'error': error_text})}\n\n".encode())
-                await response.write(b"data: [DONE]\n\n")
-                return response
-            async for line in api_response.content:
-                try:
-                    if line:
-                        if isClaudeAI:
-                            await response.write(f"data: {line.decode('utf-8')}\n\n".encode())
-                        else:
-                            try:
-                                data = json.loads(line.decode('utf-8').strip())
-                                if data.get('type') == "content_block_delta" and data.get('delta', {}).get('type') == "text_delta":
-                                    text = data['delta'].get('text', '')
-                                    message = MessageProcessor.create_chat_response(text, payload.get("model"), True)
-                                    await response.write(f"data: {json.dumps(message)}\n\n".encode())
-                            except Exception as e:
-                                logger.error(f"Error processing line: {e}")
-                except Exception as e:
-                    logger.error(f"Error processing line: {e}")
-            await response.write(b"data: [DONE]\n\n")
-    return response
-
-
-async def start_proxy_server():
-    global proxy_server_running
-    if proxy_server_running:
-        logger.info("代理服务器已运行，跳过启动")
-        return
-        
-    proxy_server_running = True
-    app = web.Application(middlewares=[auth_middleware])
-    app.router.add_post('/v1/messages', handle_message_request)
-    app.router.add_post('/v1/chat/completions', handle_message_request)
-    app.router.add_get('/v1/models', handle_models_request)
-    
-    async def health_check():
-        return web.json_response({
-            "status": "ok",
-            "message": "Zed LLM proxy is running"
-        })
-
-    app.router.add_get('/', health_check)
-
-    runner = web.AppRunner(app)
-    await runner.setup()
-    site = web.TCPSite(runner, 'localhost', CONFIG['SERVER']['PORT'])
-    await site.start()
-    logger.info(f"代理服务器启动 http://localhost:{CONFIG['SERVER']['PORT']}")
-    while True:
-        await asyncio.sleep(3600)
-
-def is_token_expiring():
-    if not token_timestamp:
-        return False
-    return (time.time() - token_timestamp) / 60 >= CONFIG['SERVER']['TOKEN_EXPIRY_WARNING_MINUTES']
-
-async def monitor_token_expiration():
-    while True:
-        await asyncio.sleep(60)
-        if is_token_expiring():
-            elapsed = int((time.time() - token_timestamp) / 60)
-            logger.warning(f"LLM令牌接近过期 (收到 {elapsed} 分钟前)")
-            if active_websocket is None:
-                logger.info("重新连接WebSocket以刷新令牌")
-                asyncio.create_task(reconnect_for_token_refresh())
-                return
-
-async def reconnect_for_token_refresh():
-    try:
-        if not CONFIG['LOGIN']['USER_ID'] or not CONFIG['LOGIN']['AUTH']:
-            logger.error("用户ID或授权令牌未设置")
-            return
-        headers = {
-            "authorization": f"{CONFIG['LOGIN']['USER_ID']} {CONFIG['LOGIN']['AUTH']}",
-            "x-zed-protocol-version": "68",
-            "x-zed-app-version": "0.178.0",
-            "x-zed-release-channel": "stable"
-        }
-        ssl_context = ssl.create_default_context()
-        ssl_context.check_hostname = False
-        ssl_context.verify_mode = ssl.CERT_NONE
-        
-        async for websocket in connect(CONFIG['API']['WS_URL'], additional_headers=headers, ssl=ssl_context):
-            try:
-                ping_task = asyncio.create_task(ping_periodically(websocket))
-                await asyncio.sleep(2)
-                await request_accept_terms_of_service(websocket)
-                await handle_messages(websocket)
-                break
-            except ConnectionClosed:
-                continue
-            except Exception as e:
-                logger.error(f"令牌刷新期间发生错误: {e}")
-                await asyncio.sleep(1)
-                continue
-            finally:
-                ping_task.cancel()
-                try:
-                    await ping_task
-                except asyncio.CancelledError:
-                    pass
-    except Exception as e:
-        logger.error(f"令牌刷新失败: {e}")
-
-async def async_main():
-    if not CONFIG['LOGIN']['USER_ID'] or not CONFIG['LOGIN']['AUTH']:
-        logger.error("用户ID或授权令牌未设置")
-        return
-    headers = {
-        "authorization": f"{CONFIG['LOGIN']['USER_ID']} {CONFIG['LOGIN']['AUTH']}",
-        "x-zed-protocol-version": "68",
-        "x-zed-app-version": "0.178.0",
-        "x-zed-release-channel": "stable"
-    }
-    ssl_context = ssl.create_default_context()
-    ssl_context.check_hostname = False
-    ssl_context.verify_mode = ssl.CERT_NONE
-    logger.info("连接到Websocket服务器")
-    async for websocket in connect(CONFIG['API']['WS_URL'], additional_headers=headers, ssl=ssl_context):
-        try:
-            ping_task = asyncio.create_task(ping_periodically(websocket))
-            token_request_task = asyncio.create_task(delayed_token_request(websocket, delay=2))
-            await handle_messages(websocket)
-            break
-        except ConnectionClosed:
-            continue
-        except Exception as e:
-            logger.error(f"意外错误: {e}")
-            await asyncio.sleep(1)
-            continue
-        finally:
-            ping_task.cancel()
-            try:
-                await ping_task
-            except asyncio.CancelledError:
-                pass
-            token_request_task.cancel()
-            try:
-                await token_request_task
-            except asyncio.CancelledError:
-                pass
-
-    while True:
-        await asyncio.sleep(3600)
-
-async def delayed_token_request(websocket, delay=2):
-    await asyncio.sleep(delay)
-    await request_accept_terms_of_service(websocket)
-
-if __name__ == "__main__":
+# Generated from trimmed zed.proto
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import runtime_version as _runtime_version
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+_runtime_version.ValidateProtobufRuntimeVersion(
+    _runtime_version.Domain.PUBLIC,
+    5,
+    29,
+    0,
+    '',
+    'zed.proto'
+)
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\tzed.proto\x12\x0czed.messages\"&\n\x06PeerId\x12\x10\n\x08owner_id\x18\x01 \x01(\r\x12\n\n\x02id\x18\x02 \x01(\r\"\xe2\x05\n\x08\x45nvelope\x12\n\n\x02id\x18\x01 \x01(\r\x12\x1a\n\rresponding_to\x18\x02 \x01(\rH\x01\x88\x01\x01\x12\x35\n\x12original_sender_id\x18\x03 \x01(\x0b\x32\x14.zed.messages.PeerIdH\x02\x88\x01\x01\x12\x14\n\x06\x61\x63k_id\x18\x8a\x02 \x01(\rH\x03\x88\x01\x01\x12$\n\x05hello\x18\x04 \x01(\x0b\x32\x13.zed.messages.HelloH\x00\x12 \n\x03\x61\x63k\x18\x05 \x01(\x0b\x32\x11.zed.messages.AckH\x00\x12$\n\x05\x65rror\x18\x06 \x01(\x0b\x32\x13.zed.messages.ErrorH\x00\x12\"\n\x04ping\x18\x07 \x01(\x0b\x32\x12.zed.messages.PingH\x00\x12\x46\n\x17\x61\x63\x63\x65pt_terms_of_service\x18\xef\x01 \x01(\x0b\x32".zed.messages.AcceptTermsOfServiceH\x00\x12W\n accept_terms_of_service_response\x18\xf0\x01 \x01(\x0b\x32*.zed.messages.AcceptTermsOfServiceResponseH\x00\x12\x33\n\rget_llm_token\x18\xeb\x01 \x01(\x0b\x32\x19.zed.messages.GetLlmTokenH\x00\x12\x44\n\x16get_llm_token_response\x18\xec\x01 \x01(\x0b\x32!.zed.messages.GetLlmTokenResponseH\x00\x42\t\n\x07payloadB\x10\n\x0e_responding_toB\x15\n\x13_original_sender_idB\t\n\x07_ack_idJ\x04\x08W\x10YJ\x06\x08\x9e\x01\x10\xa2\x01J\x06\x08\xa4\x01\x10\xa5\x01J\x06\x08\xa6\x01\x10\xaa\x01J\x06\x08\xb1\x01\x10\xba\x01J\x06\x08\xbc\x01\x10\xbd\x01J\x06\x08\xc1\x01\x10\xc4\x01J\x06\x08\xc5\x01\x10\xc6\x01J\x06\x08\xc8\x01\x10\xcb\x01J\x06\x08\xcd\x01\x10\xcf\x01J\x06\x08\xdd\x01\x10\xde\x01J\x06\x08\xe0\x01\x10\xe6\x01J\x06\x08\xf6\x01\x10\xf7\x01J\x06\x08\xf7\x01\x10\xff\x01J\x06\x08\xff\x01\x10\x81\x02\".\n\x05Hello\x12%\n\x07peer_id\x18\x01 \x01(\x0b\x32\x14.zed.messages.PeerId\"\x06\n\x04Ping\"\x05\n\x03\x41\x63k\"M\n\x05\x45rror\x12\x0f\n\x07message\x18\x01 \x01(\t\x12%\n\x04\x63ode\x18\x02 \x01(\x0e\x32\x17.zed.messages.ErrorCode\x12\x0c\n\x04tags\x18\x03 \x03(\t\"\x16\n\x14\x41\x63\x63\x65ptTermsOfService\"7\n\x1c\x41\x63\x63\x65ptTermsOfServiceResponse\x12\x17\n\x0f\x61\x63\x63\x65pted_tos_at\x18\x01 \x01(\x04\"\r\n\x0bGetLlmToken\"$\n\x13GetLlmTokenResponse\x12\r\n\x05token\x18\x01 \x01(\t*\xef\x02\n\tErrorCode\x12\x0c\n\x08Internal\x10\x00\x12\x11\n\rNoSuchChannel\x10\x01\x12\x10\n\x0c\x44isconnected\x10\x02\x12\r\n\tSignedOut\x10\x03\x12\x13\n\x0fUpgradeRequired\x10\x04\x12\r\n\tForbidden\x10\x05\x12\x0c\n\x08NeedsCla\x10\x07\x12\x13\n\x0fNotARootChannel\x10\x08\x12\x14\n\x10\x42\x61\x64PublicNesting\x10\t\x12\x13\n\x0f\x43ircularNesting\x10\n\x12\x13\n\x0fWrongMoveTarget\x10\x0b\x12\x10\n\x0cUnsharedItem\x10\x0c\x12\x11\n\rNoSuchProject\x10\r\x12$\n DevServerProjectPathDoesNotExist\x10\x10\x12\x19\n\x15RemoteUpgradeRequired\x10\x11\x12\x15\n\x11RateLimitExceeded\x10\x12\x12\x10\n\x0c\x43ommitFailed\x10\x13\"\x04\x08\x06\x10\x06\"\x04\x08\x0e\x10\x0f\x62\x06proto3')
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'zed_pb2', _globals)
+if not _descriptor._USE_C_DESCRIPTORS:
+  DESCRIPTOR._loaded_options = None
+  _globals['_ERRORCODE']._serialized_start=1085
+  _globals['_ERRORCODE']._serialized_end=1452
+  _globals['_PEERID']._serialized_start=27
+  _globals['_PEERID']._serialized_end=65
+  _globals['_ENVELOPE']._serialized_start=68
+  _globals['_ENVELOPE']._serialized_end=806
+  _globals['_HELLO']._serialized_start=808
+  _globals['_HELLO']._serialized_end=854
+  _globals['_PING']._serialized_start=856
+  _globals['_PING']._serialized_end=862
+  _globals['_ACK']._serialized_start=864
+  _globals['_ACK']._serialized_end=869
+  _globals['_ERROR']._serialized_start=871
+  _globals['_ERROR']._serialized_end=948
+  _globals['_ACCEPTTERMSOFSERVICE']._serialized_start=950
+  _globals['_ACCEPTTERMSOFSERVICE']._serialized_end=972
+  _globals['_ACCEPTTERMSOFSERVICERESPONSE']._serialized_start=974
+  _globals['_ACCEPTTERMSOFSERVICERESPONSE']._serialized_end=1029
+  _globals['_GETLLMTOKEN']._serialized_start=1031
+  _globals['_GETLLMTOKEN']._serialized_end=1044
+  _globals['_GETLLMTOKENRESPONSE']._serialized_start=1046
+  _globals['_GETLLMTOKENRESPONSE']._serialized_end=1082
+
+# Start of the actual script
+import os
+import json
+import ssl
+import time
+import asyncio
+import logging
+import aiohttp
+from aiohttp import web
+import zstandard as zstd
+from websockets.asyncio.client import connect
+from websockets.exceptions import ConnectionClosed
+import uuid
+
+from google.protobuf.json_format import MessageToDict
+
+Envelope = _sym_db.GetSymbol('zed.messages.Envelope')
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(levelname)s: %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+
+CONFIG = {
+    "API":{
+        "BASE_URL": "https://zed.dev",
+        "API_KEY": os.getenv("API_KEY","sk-123456"),
+        "BASE_API_URL": "https://collab.zed.dev",
+        "WS_URL": "wss://collab.zed.dev/rpc",
+        "LLM_API_URL": "https://llm.zed.dev/completion",
+    },
+    "LOGIN":{
+        "USER_ID": os.getenv("ZED_USER_ID"),
+        "AUTH": os.getenv("ZED_AUTH_TOKEN")
+    },
+    "SERVER":{
+        "PORT": os.getenv("PORT",5200),
+        "TOKEN_EXPIRY_WARNING_MINUTES": 50  
+    },
+    "MODELS":{
+        "claude-3-5-sonnet-20241022":"claude-3-5-sonnet-latest",
+        "claude-3-7-sonnet-20250219":"claude-3-7-sonnet-20250219"
+    }
+}
+
+highest_message_id = 0
+llm_token = None
+token_timestamp = None
+server_peer_id = None
+active_websocket = None
+proxy_server_running = False
+
+class MessageProcessor:
+    @staticmethod
+    def create_chat_response(message, model, is_stream=False):
+        base_response = {
+            "id": f"chatcmpl-{uuid.uuid4()}",
+            "created": int(time.time()),
+            "model": model
+        }
+
+        if is_stream:
+            return {
+                **base_response,
+                "object": "chat.completion.chunk",
+                "choices": [{
+                    "index": 0,
+                    "delta": {
+                        "content": message
+                    }
+                }]
+            }
+
+        return {
+            **base_response,
+            "object": "chat.completion",
+            "choices": [{
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": message
+                },
+                "finish_reason": "stop"
+            }],
+            "usage": None
+        }
+
+def decode_envelope(data):
+    try:
+        dctx = zstd.ZstdDecompressor()
+        decompressed_data = b''
+        with dctx.stream_reader(data) as reader:
+            while True:
+                chunk = reader.read(8192)
+                if not chunk:
+                    break
+                decompressed_data += chunk
+        envelope = Envelope()
+        envelope.ParseFromString(decompressed_data)
+        return MessageToDict(envelope, preserving_proto_field_name=True)
+    except Exception as e:
+        hex_preview = ' '.join(f'{byte:02x}' for byte in data[:20]) + ('...' if len(data) > 20 else '')
+        logger.error(f"无法解码消息： {e}; 数据预览： {hex_preview}")
+        return {"error": f"无法解码消息： {e}"}
+
+def compress_protobuf(data):
+    return zstd.ZstdCompressor(level=-7).compress(data)
+
+def create_message(message_type):
+    global highest_message_id
+    highest_message_id += 1
+    message_id = highest_message_id
+    envelope = Envelope(id=highest_message_id)
+
+    getattr(envelope, message_type).SetInParent()
+    return compress_protobuf(envelope.SerializeToString()), message_id
+
+async def ping_periodically(websocket):
+    while True:
+        try:
+            await websocket.ping()
+            await asyncio.sleep(1)
+        except Exception as e:
+            logger.error(f"发送ping错误: {e}")
+            break
+
+async def handle_messages(websocket):
+    global server_peer_id, llm_token, token_timestamp, active_websocket
+    active_websocket = websocket
+    try:
+        async for message in websocket:
+            message_bytes = message.encode('utf-8') if isinstance(message, str) else message
+            decoded = decode_envelope(message_bytes)
+            if "hello" in decoded:
+                server_peer_id = decoded.get('hello', {}).get('peer_id')
+            elif "accept_terms_of_service_response" in decoded:
+                await request_llm_token(websocket)
+            elif ("get_llm_token_response" in decoded and 
+                  'token' in decoded.get('get_llm_token_response', {})):
+                llm_token = decoded['get_llm_token_response']['token']
+                token_timestamp = time.time()
+                logger.info(f"LLM令牌收到  {time.ctime(token_timestamp)}")
+                if not proxy_server_running:
+                    asyncio.create_task(start_proxy_server())
+                asyncio.create_task(monitor_token_expiration())
+                logger.info("关闭WebSocket连接，直到需要刷新令牌")
+                await websocket.close()
+                active_websocket = None
+                return
+    except ConnectionClosed:
+        logger.info("连接已关闭")
+        active_websocket = None
+
+async def request_llm_token(websocket):
+    message, _ = create_message('get_llm_token')
+    logger.info("请求LLM令牌")
+    await websocket.send(message)
+
+async def request_accept_terms_of_service(websocket):
+    message, _ = create_message('accept_terms_of_service')
+    logger.info("发送同意Zed服务条款")
+    await websocket.send(message)
+
+def format_content(content):
+    if isinstance(content, str):
+        return [{"type": "text", "text": content}]
+    return content
+
+
+
+async def process_message_content(content):
+  """
+  处理消息内容，将不同类型的内容转换为字符串
+  """
+  if isinstance(content, str):
+      return content
+  
+  if isinstance(content, list):
+      return '\n'.join([item.get('text', '') for item in content])
+  
+  if isinstance(content, dict):
+      return content.get('text', None)
+  
+  return None
+
+async def transform_messages(request):
+  """
+  转换消息格式，合并系统消息并处理消息结构
+  """
+  system_message = ''  # 存储系统消息的变量
+  is_collecting_system_message = False  # 是否正在收集系统消息
+  has_processed_system_messages = False  # 是否已处理初始系统消息
+
+  converted_messages = []
+
+  for current in request.get('messages', []):
+      role = current.get('role')
+      current_content = await process_message_content(current.get('content'))
+
+      if current_content is None:
+          converted_messages.append(current)
+          continue
+
+      if role == 'system' and not has_processed_system_messages:
+          if not is_collecting_system_message:
+              # 第一次遇到system，开启收集
+              system_message = current_content
+              is_collecting_system_message = True
+          else:
+              # 继续遇到system，合并system消息
+              system_message += '\n' + current_content
+          continue
+
+      # 遇到非system消息
+      if is_collecting_system_message:
+          # 结束系统消息收集
+          is_collecting_system_message = False
+          has_processed_system_messages = True
+
+      # 如果已处理初始消息序列且再次遇到system，则转换role为user
+      if has_processed_system_messages and role == 'system':
+          role = 'user'
+
+      # 检查是否可以合并消息
+      if converted_messages and converted_messages[-1].get('role') == role:
+          converted_messages[-1]['content'][0]['text'] += '\r\n' + current_content
+      else:
+          converted_messages.append({
+              'role': role,
+              'content': [{'type': 'text', 'text': current_content}]
+          })
+
+  return {
+      'messages': converted_messages,
+      'system': system_message,
+      'model': CONFIG['MODELS'].get(request.get('model'), "claude-3-5-sonnet-latest"),
+      'max_tokens': request.get('max_tokens',8192),
+      'temperature': max(0, min(request.get('temperature', 0), 1)),
+      'top_p': max(0, min(request.get('top_p', 1), 1)),
+      'top_k': max(0, min(request.get('top_k', 0), 500)),
+      'stream': True
+  }
+
+@web.middleware
+async def auth_middleware(request, handler):
+    if CONFIG['API']['API_KEY']:
+        auth_header = request.headers.get('Authorization')
+        xapi_key_header = request.headers.get('x-api-key')
+
+        auth_password = None
+        if auth_header and auth_header.startswith('Bearer '):
+            auth_password = auth_header[7:]
+
+        if auth_password == CONFIG['API']['API_KEY'] or xapi_key_header == CONFIG['API']['API_KEY']:
+            return await handler(request)
+        else:
+            return web.json_response(
+                {"error": "Unauthorized"},
+                status=401
+            )
+
+    return await handler(request)
+
+async def handle_models_request(request):
+    return web.json_response({
+        "object": "list",
+        "data": [
+            {
+                "id": model,
+                "object": "model",
+                "created": int(time.time()),
+                "owned_by": "zed"
+            }
+            for model in CONFIG["MODELS"].keys()
+        ]
+    })
+
+async def handle_message_request(request):
+    global llm_token
+    if not llm_token:
+        return web.json_response({"error": "LLM令牌不可用"}, status=500)
+    try:
+        body = await request.json()
+        isClaudeAI = False
+        if request.path == '/v1/messages':
+            isClaudeAI = True
+            if "messages" in body:
+                for msg in body["messages"]:
+                    if "content" in msg:
+                        msg["content"] = format_content(msg["content"])
+            if "system" in body:
+                if isinstance(body["system"], list):
+                    body["system"] = "\n".join([item["text"] for item in body["system"]])
+            if "model" in body:
+                body["model"] = CONFIG['MODELS'].get(body["model"], "claude-3-5-sonnet-latest")
+        else:
+            body = await transform_messages(body)
+        with open('request_payload222.json', 'w', encoding='utf-8') as f:
+            json.dump(body, f, ensure_ascii=False, indent=2)
+        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {llm_token}"}
+        with open('ceshi.txt', 'w', encoding='utf-8') as f:
+            f.write(llm_token + '\n')
+        payload = {
+            "provider": "anthropic",
+            "model": body.get("model", "claude-3-5-sonnet-latest"),
+            "provider_request": body
+        }
+        # with open('ceshi.txt', 'w', encoding='utf-8') as f:
+        #     f.write(json.dumps(body,ensure_ascii=False) + '\n')
+        if body.get("stream", False):
+            return await handle_streaming_request(request, headers, payload, isClaudeAI)
+        else:
+            return await handle_non_streaming_request(headers, payload, isClaudeAI)
+    except Exception as e:
+        logger.error(f"处理请求时发生错误: {e}")
+        return web.json_response({"error": str(e)}, status=500)
+
+async def handle_non_streaming_request(headers, payload, isClaudeAI=False):
+    async with aiohttp.ClientSession() as session:
+        async with session.post(CONFIG['API']['LLM_API_URL'], headers=headers, json=payload) as r:
+            if r.status != 200:
+                text = await r.text()
+                logger.error(f"LLM API错误: {text}")
+                return web.json_response({"error": text}, status=r.status)
+            full_content, message_data = "", {}
+            async for line in r.content:
+                if not line:
+                    continue
+                try:
+                    event = json.loads(line.decode('utf-8').strip())
+                    et = event.get('type')
+                    if et == "message_start":
+                        message_data = event.get('message', {})
+                    elif et == "content_block_delta" and event.get('delta', {}).get('type') == "text_delta":
+                        full_content += event['delta'].get('text', '')
+                    elif et == "message_delta" and 'usage' in event:
+                        message_data['usage'] = event.get('usage')
+                    elif et == "message_stop":
+                        break
+                except Exception as e:
+                    logger.error(f"Error processing line: {e}")
+            if isClaudeAI:
+                message_data['content'] = [{"type": "text", "text": full_content}]
+            else:
+                message_data = MessageProcessor.create_chat_response(full_content, payload.get("model"), False)
+            return web.json_response(message_data)
+
+async def handle_streaming_request(request, headers, payload, isClaudeAI=False):
+    response = web.StreamResponse()
+    response.headers['Content-Type'] = 'text/event-stream'
+    response.headers['Cache-Control'] = 'no-cache'
+    response.headers['Connection'] = 'keep-alive'
+    await response.prepare(request)
+    logger.info(f"开始处理流请求")
+    async with aiohttp.ClientSession() as session:
+        async with session.post(CONFIG['API']['LLM_API_URL'], headers=headers, json=payload) as api_response:
+            if api_response.status != 200:
+                error_text = await api_response.text()
+                logger.error(f"LLM API (stream)错误: {error_text}")
+                await response.write(f"data: {json.dumps({'error': error_text})}\n\n".encode())
+                await response.write(b"data: [DONE]\n\n")
+                return response
+            async for line in api_response.content:
+                try:
+                    if line:
+                        if isClaudeAI:
+                            await response.write(f"data: {line.decode('utf-8')}\n\n".encode())
+                        else:
+                            try:
+                                data = json.loads(line.decode('utf-8').strip())
+                                if data.get('type') == "content_block_delta" and data.get('delta', {}).get('type') == "text_delta":
+                                    text = data['delta'].get('text', '')
+                                    message = MessageProcessor.create_chat_response(text, payload.get("model"), True)
+                                    await response.write(f"data: {json.dumps(message)}\n\n".encode())
+                            except Exception as e:
+                                logger.error(f"Error processing line: {e}")
+                except Exception as e:
+                    logger.error(f"Error processing line: {e}")
+            await response.write(b"data: [DONE]\n\n")
+    return response
+
+
+async def start_proxy_server():
+    global proxy_server_running
+    if proxy_server_running:
+        logger.info("代理服务器已运行，跳过启动")
+        return
+        
+    proxy_server_running = True
+    app = web.Application(middlewares=[auth_middleware])
+    app.router.add_post('/v1/messages', handle_message_request)
+    app.router.add_post('/v1/chat/completions', handle_message_request)
+    app.router.add_get('/v1/models', handle_models_request)
+    
+    async def health_check():
+        return web.json_response({
+            "status": "ok",
+            "message": "Zed LLM proxy is running"
+        })
+
+    app.router.add_get('/', health_check)
+
+    runner = web.AppRunner(app)
+    await runner.setup()
+    site = web.TCPSite(runner, 'localhost', CONFIG['SERVER']['PORT'])
+    await site.start()
+    logger.info(f"代理服务器启动 http://localhost:{CONFIG['SERVER']['PORT']}")
+    while True:
+        await asyncio.sleep(3600)
+
+def is_token_expiring():
+    if not token_timestamp:
+        return False
+    return (time.time() - token_timestamp) / 60 >= CONFIG['SERVER']['TOKEN_EXPIRY_WARNING_MINUTES']
+
+async def monitor_token_expiration():
+    while True:
+        await asyncio.sleep(60)
+        if is_token_expiring():
+            elapsed = int((time.time() - token_timestamp) / 60)
+            logger.warning(f"LLM令牌接近过期 (收到 {elapsed} 分钟前)")
+            if active_websocket is None:
+                logger.info("重新连接WebSocket���刷新令牌")
+                asyncio.create_task(reconnect_for_token_refresh())
+                return
+
+async def reconnect_for_token_refresh():
+    try:
+        if not CONFIG['LOGIN']['USER_ID'] or not CONFIG['LOGIN']['AUTH']:
+            logger.error("用户ID或授权令牌未设置")
+            return
+        headers = {
+            "authorization": f"{CONFIG['LOGIN']['USER_ID']} {CONFIG['LOGIN']['AUTH']}",
+            "x-zed-protocol-version": "68",
+            "x-zed-app-version": "0.178.0",
+            "x-zed-release-channel": "stable"
+        }
+        ssl_context = ssl.create_default_context()
+        ssl_context.check_hostname = False
+        ssl_context.verify_mode = ssl.CERT_NONE
+        
+        async for websocket in connect(CONFIG['API']['WS_URL'], additional_headers=headers, ssl=ssl_context):
+            try:
+                ping_task = asyncio.create_task(ping_periodically(websocket))
+                await asyncio.sleep(2)
+                await request_accept_terms_of_service(websocket)
+                await handle_messages(websocket)
+                break
+            except ConnectionClosed:
+                continue
+            except Exception as e:
+                logger.error(f"令牌刷新期间发生错误: {e}")
+                await asyncio.sleep(1)
+                continue
+            finally:
+                ping_task.cancel()
+                try:
+                    await ping_task
+                except asyncio.CancelledError:
+                    pass
+    except Exception as e:
+        logger.error(f"令牌刷新失败: {e}")
+
+async def async_main():
+    if not CONFIG['LOGIN']['USER_ID'] or not CONFIG['LOGIN']['AUTH']:
+        logger.error("用户ID或授权令牌未设置")
+        return
+    headers = {
+        "authorization": f"{CONFIG['LOGIN']['USER_ID']} {CONFIG['LOGIN']['AUTH']}",
+        "x-zed-protocol-version": "68",
+        "x-zed-app-version": "0.178.0",
+        "x-zed-release-channel": "stable"
+    }
+    ssl_context = ssl.create_default_context()
+    ssl_context.check_hostname = False
+    ssl_context.verify_mode = ssl.CERT_NONE
+    logger.info("连接到Websocket服务器")
+    async for websocket in connect(CONFIG['API']['WS_URL'], additional_headers=headers, ssl=ssl_context):
+        try:
+            ping_task = asyncio.create_task(ping_periodically(websocket))
+            token_request_task = asyncio.create_task(delayed_token_request(websocket, delay=2))
+            await handle_messages(websocket)
+            break
+        except ConnectionClosed:
+            continue
+        except Exception as e:
+            logger.error(f"意外错误: {e}")
+            await asyncio.sleep(1)
+            continue
+        finally:
+            ping_task.cancel()
+            try:
+                await ping_task
+            except asyncio.CancelledError:
+                pass
+            token_request_task.cancel()
+            try:
+                await token_request_task
+            except asyncio.CancelledError:
+                pass
+
+async def delayed_token_request(websocket, delay=2):
+    await asyncio.sleep(delay)
+    await request_accept_terms_of_service(websocket)
+
+if __name__ == "__main__":
     asyncio.run(async_main())