cdp-bridge 0.1.0__py3-none-any.whl

cdp_bridge/TMWebDriver.py

@@ -0,0 +1,285 @@
+import json, threading, time, uuid, queue, socket, requests, traceback
+from typing import Dict, Any, Optional, List  
+from simple_websocket_server import WebSocketServer, WebSocket  
+from bs4 import BeautifulSoup  
+import bottle, random
+from bottle import route, template, request, response
+
+class Session:
+    def __init__(self, session_id, info, client=None):
+        self.id = session_id
+        self.info = info
+        self.connect_at = time.time()
+        self.disconnect_at = None
+        self.type = info.get('type', 'ws')
+        self.ws_client = client if self.type in ('ws', 'ext_ws') else None
+        self.http_queue = client if self.type == 'http' else None
+    @property
+    def url(self): return self.info.get('url', '')
+    def is_active(self):
+        if self.type == 'http' and time.time() - self.connect_at > 60: self.mark_disconnected()
+        return self.disconnect_at is None
+    def reconnect(self, client, info):
+        self.info = info
+        self.type = info.get('type', 'ws')
+        if self.type in ('ws', 'ext_ws'):
+            self.ws_client = client
+            self.http_queue = None
+        elif self.type == 'http':
+            self.http_queue = client
+        self.connect_at = time.time()
+        self.disconnect_at = None
+    def mark_disconnected(self):
+        if self.is_active(): print(f"Tab disconnected: {self.url} (Session: {self.id})")
+        self.disconnect_at = time.time()
+
+
+class TMWebDriver:  
+    def __init__(self, host: str = '127.0.0.1', port: int = 18765):  
+        self.host, self.port = host, port
+        self.sessions, self.results, self.acks = {}, {}, {}
+        self.default_session_id = None  
+        self.latest_session_id = None  
+        self.is_remote = socket.socket().connect_ex((host, port+1)) == 0
+        if not self.is_remote:  
+            self.start_ws_server()  
+            self.start_http_server()
+        else:
+            self.remote = f'http://{self.host}:{self.port+1}/link'
+
+    def start_http_server(self):
+        self.app = app = bottle.Bottle()
+
+        @app.route('/api/longpoll', method=['GET', 'POST'])
+        def long_poll():
+            data = request.json
+            session_id = data.get('sessionId')  
+            session_info = {'url': data.get('url'), 'title': data.get('title', ''), 'type': 'http'}  
+            if session_id not in self.sessions: 
+                session = Session(session_id, session_info, queue.Queue())
+                print(f"Browser http connected: {session.url} (Session: {session_id})")  
+                self.sessions[session_id] = session
+            session = self.sessions[session_id]
+            if session.disconnect_at is not None and session.type != 'http': session.reconnect(queue.Queue(), session_info)
+            session.disconnect_at = None
+            if session.type == 'http': msgQ = session.http_queue
+            else: return json.dumps({"id": "", "ret": "use ws"})
+            session.connect_at = start_time = time.time()
+            while time.time() - start_time < 5:
+                try:
+                    msg = msgQ.get(timeout=0.2)
+                    try: self.acks[json.loads(msg).get('id','')] = True
+                    except: traceback.print_exc()
+                    return msg
+                except queue.Empty: continue
+            return json.dumps({"id": "", "ret": "next long-poll"})
+
+        @app.route('/api/result', method=['GET','POST'])
+        def result():
+            data = request.json
+            if data.get('type') == 'result':  
+                self.results[data.get('id')] = {'success': True, 'data': data.get('result'), 'newTabs': data.get('newTabs', [])}  
+            elif data.get('type') == 'error':  
+                self.results[data.get('id')] = {'success': False, 'data': data.get('error'), 'newTabs': data.get('newTabs', [])}  
+            return 'ok'
+
+        @app.route('/link', method=['GET','POST'])
+        def link():
+            data = request.json
+            if data.get('cmd') == 'get_all_sessions': return json.dumps({'r': self.get_all_sessions()}, ensure_ascii=False)  
+            if data.get('cmd') == 'find_session': 
+                url_pattern = data.get('url_pattern', '')
+                return json.dumps({'r': self.find_session(url_pattern)}, ensure_ascii=False)
+            if data.get('cmd') == 'execute_js':
+                session_id = data.get('sessionId')
+                code = data.get('code')
+                timeout = float(data.get('timeout', 10.0))
+                try:
+                    result = self.execute_js(code, timeout=timeout, session_id=session_id)
+                    print('[remote result]', (str(code)[:50] + ' RESULT:' +str(result)[:50]).replace('\n', ' '))
+                    return json.dumps({'r': result}, ensure_ascii=False)
+                except Exception as e:
+                    return json.dumps({'r': {'error': str(e)}}, ensure_ascii=False)
+            return 'ok'
+        def run():
+            from wsgiref.simple_server import make_server, WSGIServer, WSGIRequestHandler
+            from socketserver import ThreadingMixIn
+            class _T(ThreadingMixIn, WSGIServer): pass
+            class _H(WSGIRequestHandler):
+                def log_request(self, *a): pass
+            make_server(self.host, self.port+1, app, server_class=_T, handler_class=_H).serve_forever()
+        http_thread = threading.Thread(target=run, daemon=True)
+        http_thread.start()  
+
+    def clean_sessions(self):
+        sids = list(self.sessions.keys())
+        for sid in sids:
+            session = self.sessions[sid]
+            if not session.is_active() and time.time() - session.disconnect_at > 600:
+                del self.sessions[sid]
+    
+    def start_ws_server(self) -> None:  
+        driver = self  
+        class JSExecutor(WebSocket):  
+            def handle(self) -> None:  
+                try:  
+                    data = json.loads(self.data)  
+                    if data.get('type') == 'ready':  
+                        session_id = data.get('sessionId')  
+                        session_info = {'url': data.get('url'), 'title': data.get('title', ''),
+                            'connected_at': time.time(), 'type': 'ws'}  
+                        driver._register_client(session_id, self, session_info)  
+                    elif data.get('type') in ['ext_ready', 'tabs_update']:
+                        tabs = data.get('tabs', [])
+                        current_tab_ids = {str(tab['id']) for tab in tabs}
+                        print(f"Received tabs update: {current_tab_ids}")
+                        for sid in list(driver.sessions.keys()):
+                            sess = driver.sessions[sid]
+                            if sess.type == 'ext_ws' and sid not in current_tab_ids:
+                                sess.mark_disconnected()
+                        for tab in tabs:
+                            session_id = str(tab['id'])
+                            session_info = {'url': tab.get('url'), 'title': tab.get('title', ''), 'connected_at': time.time(), 'type': 'ext_ws'}
+                            sess = driver.sessions.get(session_id)
+                            if sess and sess.is_active(): sess.info = session_info
+                            else: driver._register_client(session_id, self, session_info)
+                    elif data.get('type') == 'ack': driver.acks[data.get('id','')] = True
+                    elif data.get('type') == 'result':  
+                        driver.results[data.get('id')] = {'success': True, 'data': data.get('result'), 'newTabs': data.get('newTabs', [])}  
+                    elif data.get('type') == 'error':  
+                        driver.results[data.get('id')] = {'success': False, 'data': data.get('error'), 'newTabs': data.get('newTabs', [])}  
+                except Exception as e:  
+                    print(f"Error handling message: {e}")  
+                    if hasattr(self, 'data'): print(self.data)  
+            def connected(self): (f"New connection from {self.address}")  
+            def handle_close(self): 
+                print(f"WS Connection closed: {self.address}")
+                driver._unregister_client(self)  
+        
+        self.server = WebSocketServer(self.host, self.port, JSExecutor)  
+        server_thread = threading.Thread(target=self.server.serve_forever)  
+        server_thread.daemon = True  
+        server_thread.start()  
+        print(f"WebSocket server running on ws://{self.host}:{self.port}")  
+    
+    def _register_client(self, session_id: str, client: WebSocket, session_info) -> None:  
+        is_new_session = session_id not in self.sessions
+
+        if is_new_session:
+            session = Session(session_id, session_info, client)
+            self.sessions[session_id] = session            
+            print(f"New tab connected: {session.url} (Session: {session_id})")  
+        else:
+            session = self.sessions[session_id]
+            session.reconnect(client, session_info)
+            print(f"Tab reconnected: {session.url} (Session: {session_id})")  
+
+        self.latest_session_id = session_id
+        if self.default_session_id is None: self.default_session_id = session_id 
+    
+    def _unregister_client(self, client: WebSocket) -> None:  
+        for session in self.sessions.values():
+            if session.ws_client == client: session.mark_disconnected()
+    
+    def execute_js(self, code, timeout=15, session_id=None) -> Any:  
+        if session_id is None: session_id = self.default_session_id  
+        if self.is_remote:
+            print('remote_execute_js')
+            response = self._remote_cmd({"cmd": "execute_js", "sessionId": session_id, 
+                                         "code": code, "timeout": str(timeout)}).get('r', {})
+            if response.get('error'): raise Exception(response['error'])
+            return response
+ 
+        session = self.sessions.get(session_id)
+        if not session or not session.is_active(): 
+            time.sleep(3)
+            session = self.sessions.get(session_id)
+            if not session or not session.is_active(): 
+                alive_sessions = [s for s in self.sessions.values() if s.is_active()]
+                if alive_sessions:
+                    session = alive_sessions[0]  
+                    print(f"会话 {session_id} 未连接，自动切换到最新活动会话: {session.id}")
+                    session_id = self.default_session_id = session.id
+                if not session or not session.is_active(): 
+                    raise ValueError(f"会话ID {session_id} 未连接")  
+
+        tp = session.type
+        if tp not in ('ws', 'http', 'ext_ws'):
+            raise ValueError(f"Unsupported session type: {tp}")
+        exec_id = str(uuid.uuid4())  
+        payload_dict = {'id': exec_id, 'code': code}
+        if tp == 'ext_ws': payload_dict['tabId'] = int(session.id)
+        payload = json.dumps(payload_dict)
+
+        if tp in ['ws', 'ext_ws']: session.ws_client.send_message(payload)  
+        elif tp == 'http': session.http_queue.put(payload)
+
+        start_time = time.time()  
+        self.clean_sessions() 
+        hasjump = acked = False
+
+        while exec_id not in self.results:  
+            time.sleep(0.2)  
+            if not acked and exec_id in self.acks:
+                acked = True; start_time = time.time()
+            if tp in ['ws', 'ext_ws']:
+                if not session.is_active(): hasjump = True
+                if hasjump and session.is_active():
+                    return {'result': f"Session {session_id} reloaded.", "closed":1}
+            if time.time() - start_time > timeout:  
+                if tp in ['ws', 'ext_ws']:
+                    if hasjump: return {'result': f"Session {session_id} reloaded and new page is loading...", 'closed':1}
+                    if acked: return {"result": f"No response data in {timeout}s (ACK received, script may still be running)"}
+                    return {"result": f"No response data in {timeout}s (no ACK, script may not have been delivered)"}
+                elif tp == 'http':
+                    if acked: return {"result": f"Session {session_id} no response in {timeout}s (delivered but no result)"}
+                    return {"result": f"Session {session_id} no response in {timeout}s (script not polled)"}
+        
+        result = self.results.pop(exec_id)  
+        if exec_id in self.acks: self.acks.pop(exec_id)
+        if not result['success']: raise Exception(result['data'])  
+        rr = {'data': result['data']}
+        newtabs = result.get('newTabs', []); [x.pop('ts', None) for x in newtabs]
+        if newtabs: rr['newTabs'] = newtabs
+        return rr
+    
+    def _remote_cmd(self, cmd):
+        try: return requests.post(self.remote, headers={"Content-Type": "application/json"}, json=cmd, timeout=30).json()
+        except (ConnectionError, requests.exceptions.ConnectionError):
+            raise ConnectionError("TMWebDriver master未运行，看tmwebdriver_sop启动master")
+
+    def get_all_sessions(self):  
+        if self.is_remote:
+            return self._remote_cmd({"cmd": "get_all_sessions"}).get('r', [])
+        return [{'id': session.id, **session.info} for session in self.sessions.values()
+                if session.is_active()]  
+
+    def get_session_dict(self):
+        return {session['id']: session['url'] for session in self.get_all_sessions()}
+        
+    def find_session(self, url_pattern: str):
+        if url_pattern == '': 
+            session = self.sessions.get(self.latest_session_id)
+            return [(session.id, session.info)] if session else []
+        matching_sessions = []  
+        for session in self.sessions.values():
+            if not session.is_active(): continue
+            if 'url' in session.info and url_pattern in session.info['url']:  
+                matching_sessions.append((session.id, session.info))  
+        return matching_sessions
+
+    def set_session(self, url_pattern: str) -> bool:  
+        if self.is_remote:
+            matched = self._remote_cmd({"cmd": "find_session", "url_pattern": url_pattern}).get('r', [])
+        else:
+            matched = self.find_session(url_pattern)
+        if not matched: return print(f"警告: 未找到URL包含 '{url_pattern}' 的会话")  
+        if len(matched) > 1: print(f"警告: 找到多个URL包含 '{url_pattern}' 的会话，选择第一个")  
+        self.default_session_id, info = matched[0]
+        print(f"成功设置默认会话: {self.default_session_id}: {info['url']}")  
+        return self.default_session_id  
+    
+    def jump(self, url, timeout=10): self.execute_js(f"window.location.href='{url}'", timeout=timeout)
+    
+if __name__ == "__main__":
+    driver = TMWebDriver(host='127.0.0.1', port=18765)

cdp_bridge/__init__.py

@@ -0,0 +1,23 @@
+import argparse
+from importlib import resources
+
+from .server import mcp
+
+
+def main():
+    """Run the CDP Bridge MCP server."""
+    parser = argparse.ArgumentParser(
+        description="Run the CDP Bridge MCP server for browser automation through the companion extension."
+    )
+    parser.parse_args()
+    mcp.run()
+
+
+def extension_path():
+    """Print the packaged Chrome extension directory."""
+    extension_dir = resources.files(__package__) / "tmwd_cdp_bridge"
+    print(extension_dir)
+
+
+if __name__ == "__main__":
+    main()

cdp_bridge/__main__.py

@@ -0,0 +1,5 @@
+from . import main
+
+
+if __name__ == "__main__":
+    main()

cdp_bridge/server.py

@@ -0,0 +1,164 @@
+import asyncio, json
+import importlib
+
+from mcp.server.fastmcp import FastMCP
+
+from . import simphtml
+
+mcp = FastMCP("tmwebdriver-bridge")
+
+driver = None
+
+def get_driver():
+    global driver
+    if driver is None:
+        from .TMWebDriver import TMWebDriver
+        driver = TMWebDriver()
+    return driver
+
+
+@mcp.tool()
+async def browser_get_tabs() -> str:
+    """Get all open browser tabs with their IDs, URLs, and titles."""
+    def _run():
+        d = get_driver()
+        sessions = d.get_all_sessions()
+        for s in sessions:
+            s.pop('connected_at', None)
+            s.pop('type', None)
+        return json.dumps({"tabs": sessions, "active_tab": d.default_session_id}, ensure_ascii=False)
+    return await asyncio.to_thread(_run)
+
+
+@mcp.tool()
+async def browser_scan(tabs_only: bool = False, switch_tab_id: str = "", text_only: bool = False) -> str:
+    """Get simplified HTML content of the active tab plus tab list. The HTML is optimized for LLM consumption (stripped of scripts, styles, invisible elements).
+
+    Args:
+        tabs_only: Only return tab list without page content (saves tokens).
+        switch_tab_id: Switch to this tab before scanning.
+        text_only: Return plain text instead of simplified HTML.
+    """
+    def _run():
+        d = get_driver()
+        if len(d.get_all_sessions()) == 0:
+            return json.dumps({"status": "error", "msg": "No browser tabs connected. Ensure Chrome extension is running."}, ensure_ascii=False)
+
+        if switch_tab_id:
+            d.default_session_id = switch_tab_id
+
+        tabs = []
+        for sess in d.get_all_sessions():
+            sess.pop('connected_at', None)
+            sess.pop('type', None)
+            sess['url'] = sess.get('url', '')[:80]
+            tabs.append(sess)
+
+        result = {
+            "status": "success",
+            "metadata": {"tabs_count": len(tabs), "tabs": tabs, "active_tab": d.default_session_id}
+        }
+        if not tabs_only:
+            importlib.reload(simphtml)
+            result["content"] = simphtml.get_html(d, cutlist=True, maxchars=35000, text_only=text_only)
+        return json.dumps(result, ensure_ascii=False, default=str)
+    return await asyncio.to_thread(_run)
+
+
+@mcp.tool()
+async def browser_execute_js(script: str, switch_tab_id: str = "", no_monitor: bool = False) -> str:
+    """Execute JavaScript in the browser and capture results plus DOM changes.
+
+    Args:
+        script: JavaScript code to execute (or JSON command for CDP operations).
+        switch_tab_id: Switch to this tab before executing.
+        no_monitor: Skip DOM change monitoring (faster, less info).
+    """
+    def _run():
+        d = get_driver()
+        if len(d.get_all_sessions()) == 0:
+            return json.dumps({"status": "error", "msg": "No browser tabs connected."}, ensure_ascii=False)
+        if switch_tab_id:
+            d.default_session_id = switch_tab_id
+        importlib.reload(simphtml)
+        result = simphtml.execute_js_rich(script, d, no_monitor=no_monitor)
+        return json.dumps(result, ensure_ascii=False, default=str)
+    return await asyncio.to_thread(_run)
+
+
+@mcp.tool()
+async def browser_switch_tab(tab_id: str) -> str:
+    """Switch the active browser tab.
+
+    Args:
+        tab_id: The tab ID to switch to (from browser_get_tabs).
+    """
+    def _run():
+        d = get_driver()
+        d.default_session_id = tab_id
+        session = d.sessions.get(tab_id)
+        if session and session.is_active():
+            return json.dumps({"status": "success", "active_tab": tab_id, "url": session.info.get('url', '')}, ensure_ascii=False)
+        return json.dumps({"status": "error", "msg": f"Tab {tab_id} not found or disconnected."}, ensure_ascii=False)
+    return await asyncio.to_thread(_run)
+
+
+@mcp.tool()
+async def browser_navigate(url: str) -> str:
+    """Navigate the active tab to a URL.
+
+    Args:
+        url: The URL to navigate to.
+    """
+    def _run():
+        d = get_driver()
+        if len(d.get_all_sessions()) == 0:
+            return json.dumps({"status": "error", "msg": "No browser tabs connected."}, ensure_ascii=False)
+        d.jump(url, timeout=10)
+        return json.dumps({"status": "success", "msg": f"Navigating to {url}"}, ensure_ascii=False)
+    return await asyncio.to_thread(_run)
+
+
+@mcp.tool()
+async def browser_screenshot(tab_id: str = "") -> str:
+    """Take a screenshot of the active tab (returns base64 PNG).
+
+    Args:
+        tab_id: Optional tab ID to screenshot. Uses active tab if empty.
+    """
+    def _run():
+        d = get_driver()
+        if len(d.get_all_sessions()) == 0:
+            return json.dumps({"status": "error", "msg": "No browser tabs connected."}, ensure_ascii=False)
+        cmd = {"cmd": "cdp", "method": "Page.captureScreenshot", "params": {"format": "png"}}
+        if tab_id:
+            cmd["tabId"] = int(tab_id)
+        result = d.execute_js(json.dumps(cmd))
+        data = result.get('data', {})
+        if isinstance(data, dict) and 'data' in data:
+            return json.dumps({"status": "success", "format": "png", "base64": data['data']}, ensure_ascii=False)
+        return json.dumps({"status": "success", "data": data}, ensure_ascii=False, default=str)
+    return await asyncio.to_thread(_run)
+
+
+@mcp.tool()
+async def browser_cookies(url: str = "") -> str:
+    """Get cookies for the current page or a specific URL.
+
+    Args:
+        url: URL to get cookies for. If empty, gets cookies for the active tab's URL.
+    """
+    def _run():
+        d = get_driver()
+        if len(d.get_all_sessions()) == 0:
+            return json.dumps({"status": "error", "msg": "No browser tabs connected."}, ensure_ascii=False)
+        cmd = {"cmd": "cookies"}
+        if url:
+            cmd["url"] = url
+        result = d.execute_js(json.dumps(cmd))
+        return json.dumps({"status": "success", "cookies": result.get('data', [])}, ensure_ascii=False, default=str)
+    return await asyncio.to_thread(_run)
+
+
+if __name__ == "__main__":
+    mcp.run()