Create app.py

app.py

@@ -0,0 +1,461 @@
+# 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 random
+import base64
+import json
+import urllib.parse
+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 rsa
+
+from google.protobuf.json_format import MessageToDict
+
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(levelname)s: %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+BASE_URL = "https://zed.dev"
+BASE_API_URL = "https://collab.zed.dev"
+WS_URL = "wss://collab.zed.dev/rpc"
+LLM_API_URL = "https://llm.zed.dev/completion"
+PROXY_PORT = 9090
+TOKEN_EXPIRY_WARNING_MINUTES = 50
+
+highest_message_id = 0
+llm_token = None
+token_timestamp = None
+server_peer_id = None
+active_websocket = None
+proxy_server_running = False
+
+def generate_signin_url():
+    port = random.randint(30000, 55000)
+    (pubkey, privkey) = rsa.newkeys(2048, exponent=65537)
+    der_public_key = rsa.PublicKey.save_pkcs1(pubkey, format='DER')
+    b64_public_key = base64.urlsafe_b64encode(der_public_key).decode('utf-8')
+    signin_url = f"{BASE_URL}/native_app_signin?native_app_port={port}&native_app_public_key={b64_public_key}"
+    return signin_url, privkey
+
+def decrypt_access_token(access_token_b64, private_key):
+    encrypted_token = base64.urlsafe_b64decode(access_token_b64)
+    try:
+        decrypted = rsa.decrypt(encrypted_token, private_key)
+        return decrypted.decode('utf-8')
+    except Exception as e:
+        try:
+            decrypted = rsa.decrypt(encrypted_token, private_key, 'PKCS1_v1_5')
+            return decrypted.decode('utf-8')
+        except Exception as inner_e:
+            raise Exception(f"Decryption failed (OAEP -> PKCS#1): {e}, {inner_e}")
+
+def do_auth_flow():
+    signin_url, private_key = generate_signin_url()
+    print("\n!Sign in with GitHub (wait a few seconds after signing in for the redirect):\n", signin_url)
+    print("\nPaste the full callback URL where the browser redirected you here:")
+    callback_url = input().strip()
+
+    parsed_url = urllib.parse.urlparse(callback_url)
+    query_params = urllib.parse.parse_qs(parsed_url.query)
+
+    user_id = query_params.get('user_id', [''])[0]
+    access_token_b64 = query_params.get('access_token', [''])[0]
+
+    if not user_id or not access_token_b64:
+        raise Exception("Missing user_id or access_token in callback URL.")
+
+    decrypted_token = decrypt_access_token(access_token_b64, private_key)
+    login_data = {"user_id": user_id, "auth": decrypted_token}
+    with open('login.json', 'w') as f:
+        json.dump(login_data, f, indent=2)
+    logger.info("Success! Authentication data saved to login.json")
+    logger.info(f"User ID: {user_id}")
+
+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"Unable to decode message: {e}; data preview: {hex_preview}")
+        return {"error": f"Unable to decode message: {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"Error sending 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 token received at {time.ctime(token_timestamp)}")
+                if not proxy_server_running:
+                    asyncio.create_task(start_proxy_server())
+                asyncio.create_task(monitor_token_expiration())
+                logger.info("Closing WebSocket connection until token refresh is needed")
+                await websocket.close()
+                active_websocket = None
+                return
+    except ConnectionClosed:
+        logger.info("Connection closed")
+        active_websocket = None
+
+async def request_llm_token(websocket):
+    message, _ = create_message('get_llm_token')
+    logger.info("Requesting the LLM token")
+    await websocket.send(message)
+
+async def request_accept_terms_of_service(websocket):
+    message, _ = create_message('accept_terms_of_service')
+    logger.info("Sending consent for the Zed Terms of Service")
+    await websocket.send(message)
+
+def format_content(content):
+    if isinstance(content, str):
+        return [{"type": "text", "text": content}]
+    return content
+
+async def handle_message_request(request):
+    global llm_token
+    if not llm_token:
+        return web.json_response({"error": "LLM token not available"}, status=500)
+    try:
+        body = await request.json()
+        # Zed's Anthropic API weirdly doesn't want flat strings for messages
+        if "messages" in body:
+            for msg in body["messages"]:
+                if "content" in msg:
+                    msg["content"] = format_content(msg["content"])
+        # But wants flat strings for system messages
+        if "system" in body:
+            if isinstance(body["system"], list):
+                body["system"] = "\n".join([item["text"] for item in body["system"]])
+            
+        headers = {"Content-Type": "application/json", "Authorization": f"Bearer {llm_token}"}
+        payload = {
+            "provider": "anthropic",
+            "model": body.get("model", "claude-3-5-sonnet"),
+            "provider_request": body
+        }
+        if body.get("stream", False):
+            return await handle_streaming_request(request, headers, payload)
+        else:
+            return await handle_non_streaming_request(headers, payload)
+    except Exception as e:
+        logger.error(f"Error processing request: {e}")
+        return web.json_response({"error": str(e)}, status=500)
+
+async def handle_non_streaming_request(headers, payload):
+    async with aiohttp.ClientSession() as session:
+        async with session.post(LLM_API_URL, headers=headers, json=payload) as r:
+            if r.status != 200:
+                text = await r.text()
+                logger.error(f"LLM API error: {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}")
+            message_data['content'] = [{"type": "text", "text": full_content}]
+            return web.json_response(message_data)
+
+async def handle_streaming_request(request, headers, payload):
+    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)
+    async with aiohttp.ClientSession() as session:
+        async with session.post(LLM_API_URL, headers=headers, json=payload, timeout=60) as api_response:
+            if api_response.status != 200:
+                error_text = await api_response.text()
+                logger.error(f"LLM API (stream) error: {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:
+                if line:
+                    await response.write(f"data: {line.decode('utf-8')}\n\n".encode())
+            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("Proxy server already running, skipping startup")
+        return
+        
+    proxy_server_running = True
+    app = web.Application()
+    
+    app.router.add_post('/v1/messages', handle_message_request)
+    runner = web.AppRunner(app)
+    await runner.setup()
+    site = web.TCPSite(runner, 'localhost', PROXY_PORT)
+    await site.start()
+    logger.info(f"Proxy server started at http://localhost:{PROXY_PORT}")
+
+    # Start localtunnel
+    try:
+        import subprocess
+        process = subprocess.Popen(
+            ['lt', '--port', str(PROXY_PORT)],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            universal_newlines=True
+        )
+        
+        # Wait for the tunnel URL
+        while True:
+            output = process.stdout.readline()
+            if output:
+                if 'your url is:' in output.lower():
+                    url = output.split(':')[-1].strip()
+                    logger.info(f"\nTunnel established!")
+                    logger.info(f"Public URL: {url}")
+                    logger.info("\nKeeping tunnel alive. Press Ctrl+C to stop.")
+                    break
+            
+            # Check if process failed
+            if process.poll() is not None:
+                error = process.stderr.read()
+                logger.error(f"Error starting tunnel: {error}")
+                break
+            
+            await asyncio.sleep(0.1)
+        
+        while True:
+            await asyncio.sleep(3600)
+            
+    except Exception as e:
+        logger.error(f"Error setting up tunnel: {e}")
+        while True:
+            await asyncio.sleep(3600)
+
+def is_token_expiring():
+    if not token_timestamp:
+        return False
+    return (time.time() - token_timestamp) / 60 >= 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 token is approaching expiration (received {elapsed} minutes ago)")
+            if active_websocket is None:
+                logger.info("Reconnecting WebSocket for token refresh")
+                asyncio.create_task(reconnect_for_token_refresh())
+                return
+
+async def reconnect_for_token_refresh():
+    try:
+        with open("login.json") as f:
+            login_data = json.load(f)
+        headers = {
+            "authorization": login_data["user_id"] + " " + login_data["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(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"Error during token refresh: {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"Failed to reconnect for token refresh: {e}")
+
+async def async_main():
+    with open("login.json") as f:
+        login_data = json.load(f)
+    headers = {
+        "authorization": login_data["user_id"] + " " + login_data["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("Connecting to the WebSocket server")
+    async for websocket in connect(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"Unexpected error: {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)
+
+def main():
+    need_auth = False
+    if not os.path.exists("login.json"):
+        need_auth = True
+    else:
+        try:
+            with open("login.json") as f:
+                data = json.load(f)
+                if not data.get("user_id") or not data.get("auth"):
+                    need_auth = True
+        except:
+            need_auth = True
+
+    if need_auth:
+        do_auth_flow()
+    else:
+        logger.info("login.json seems to be present, skipping authentication flow")
+    try:
+        asyncio.run(async_main())
+    except KeyboardInterrupt:
+        pass
+
+if __name__ == "__main__":
+    main()