Просмотр исходного кода

update: 支持账号登录后存储cookie&自动登录

guantao 3 недель назад
Родитель
Сommit
1473f022d0
7 измененных файлов с 132 добавлено и 71 удалено
  1. 1 1
      .gitignore
  2. 2 0
      agent/core/runner.py
  3. 1 1
      agent/memory/skills/core.md
  4. 64 68
      agent/tools/builtin/browser/baseClass.py
  5. 1 1
      examples/research/test.prompt
  6. 30 0
      export.py
  7. 33 0
      login.py

+ 1 - 1
.gitignore

@@ -52,7 +52,7 @@ Thumbs.db
 .env
 debug.log
 info.log
-.browser_use_files
+.cache
 output
 
 

+ 2 - 0
agent/core/runner.py

@@ -120,6 +120,8 @@ BUILTIN_TOOLS = [
     "browser_ensure_login_with_cookies",
     "browser_wait_for_user_action",
     "browser_done",
+    "browser_export_cookies",
+    "browser_load_cookies"
 ]
 
 

+ 1 - 1
agent/memory/skills/core.md

@@ -88,7 +88,7 @@ goal(abandon="方案A需要Redis,环境没有")
 - **必须先获取索引**: 所有 `index` 参数都需要先通过 `browser_get_selector_map` 获取
 - **高级工具**:优先使用`browser_extract_content`, `browser_read_long_content`等工具获取数据,而不是使用`browser_get_selector_map`获取索引后手动解析
 - **操作后等待**: 任何可能触发页面变化的操作(点击、输入、滚动)后都要调用 `browser_wait`
-- **登录处理**: 需要登录的网站使用 `browser_ensure_login_with_cookies(cookie_type="xhs")` 注入Cookie
+- **登录处理**: 需要登录的网站使用 `browser_ensure_login_with_cookies` 注入Cookie
 - **复杂操作用JS**: 当标准工具无法满足时,使用 `browser_evaluate` 执行JavaScript代码
 
 ### 工具分类

+ 64 - 68
agent/tools/builtin/browser/baseClass.py

@@ -37,7 +37,7 @@ Native Browser-Use Tools Adapter
 3. 任务结束时调用 cleanup_browser_session()
 
 文件操作说明:
-- 浏览器专用文件目录:.browser_use_files/ (在当前工作目录下)
+- 浏览器专用文件目录:.cache/.browser_use_files/ (在当前工作目录下)
   用于存储浏览器会话产生的临时文件(下载、上传、截图等)
 - 一般文件操作:请使用 agent.tools.builtin 中的文件工具 (read_file, write_file, edit_file)
   这些工具功能更完善,支持diff预览、智能匹配、分页读取等
@@ -367,7 +367,7 @@ async def init_browser_session(
     # 创建文件系统实例(用于浏览器会话产生的文件)
     # 注意:这个目录仅用于浏览器操作相关的临时文件(下载、上传、截图等)
     # 对于一般文件读写操作,请使用 agent.tools.builtin 中的文件工具
-    base_dir = Path.cwd() / ".browser_use_files"
+    base_dir = Path.cwd() / ".cache/.browser_use_files"
     base_dir.mkdir(parents=True, exist_ok=True)
     _file_system = FileSystem(base_dir=str(base_dir))
 
@@ -1421,7 +1421,7 @@ async def _detect_and_download_pdf_via_cdp(browser) -> Optional[str]:
         pdf_bytes = base64.b64decode(base64_data)
 
         # 保存到本地
-        save_dir = Path.cwd() / ".browser_use_files"
+        save_dir = Path.cwd() / ".cache/.browser_use_files"
         save_dir.mkdir(parents=True, exist_ok=True)
 
         filename = Path(parsed.path).name if parsed.path else ""
@@ -1887,110 +1887,106 @@ async def browser_done(text: str, success: bool = True,
 # Cookie 持久化工具
 # ============================================================
 
-_COOKIES_DIR = Path(__file__).parent.parent.parent.parent / ".cookies"
+_COOKIES_DIR = Path(__file__).parent.parent.parent.parent.parent / ".cache/.cookies"
 
 @tool()
-async def browser_export_cookies(name: str = "") -> ToolResult:
+async def browser_export_cookies(name: str = "", account: str = "") -> ToolResult:
     """
-    导出当前浏览器的所有 Cookie 到本地 JSON 文件。
-    Export all browser cookies to a local JSON file.
-
-    登录成功后调用此工具,下次启动时可通过 browser_load_cookies 恢复登录态。
+    导出当前浏览器的所有 Cookie 到本地 .cookies/ 目录。
+    文件命名格式:{域名}_{账号名}.json,如 bilibili.com_zhangsan.json
+    登录成功后调用此工具,下次可通过 browser_load_cookies 恢复登录态。
 
     Args:
-        name: 保存名称(可选,默认用当前域名)
-
-    Returns:
-        ToolResult: 导出结果
+        name: 自定义文件名(可选,提供则忽略自动命名)
+        account: 账号名称(可选,用于区分同一网站的不同账号)
     """
     try:
         browser, _ = await get_browser_session()
-        cookies = await browser._cdp_get_cookies()
 
-        if not cookies:
-            return ToolResult(
-                title="Cookie 导出",
-                output="当前浏览器没有 Cookie",
-                long_term_memory="导出 Cookie 失败:无 Cookie"
-            )
+        # 获取所有 Cookie(CDP 格式)
+        all_cookies = await browser._cdp_get_cookies()
+        if not all_cookies:
+            return ToolResult(title="Cookie 导出", output="当前浏览器没有 Cookie", long_term_memory="无 Cookie 可导出")
+
+        # 获取当前域名,用于过滤和命名
+        from urllib.parse import urlparse
+        current_url = await browser.get_current_page_url() or ''
+        domain = urlparse(current_url).netloc.replace("www.", "") or "default"
 
-        # 确定文件名
         if not name:
-            url = getattr(browser, 'current_url', '') or ''
-            from urllib.parse import urlparse
-            parsed = urlparse(url)
-            name = parsed.netloc.replace("www.", "") or "default"
+            name = f"{domain}_{account}" if account else domain
 
+        # 只保留当前域名的 cookie(过滤第三方)
+        cookies = [c for c in all_cookies if domain in c.get("domain", "").lstrip(".")]
+
+        # 保存
         _COOKIES_DIR.mkdir(parents=True, exist_ok=True)
         cookie_file = _COOKIES_DIR / f"{name}.json"
-
-        cookie_file.write_text(
-            json.dumps(cookies, ensure_ascii=False, indent=2),
-            encoding="utf-8"
-        )
+        cookie_file.write_text(json.dumps(cookies, ensure_ascii=False, indent=2), encoding="utf-8")
 
         return ToolResult(
             title="Cookie 已导出",
-            output=f"已保存 {len(cookies)} 条 Cookie 到 {cookie_file.name}",
-            long_term_memory=f"导出 {len(cookies)} 条 Cookie 到 .cookies/{cookie_file.name}"
+            output=f"已保存 {len(cookies)} 条 Cookie 到 .cookies/{name}.json(从 {len(all_cookies)} 条中过滤当前域名)",
+            long_term_memory=f"导出 {len(cookies)} 条 Cookie 到 .cookies/{name}.json"
         )
-
     except Exception as e:
-        return ToolResult(
-            title="Cookie 导出失败",
-            output="",
-            error=f"Failed to export cookies: {str(e)}",
-            long_term_memory="导出 Cookie 失败"
-        )
+        return ToolResult(title="Cookie 导出失败", output="", error=str(e), long_term_memory="导出 Cookie 失败")
 
 
 @tool()
-async def browser_load_cookies(name: str, url: str = "") -> ToolResult:
+async def browser_load_cookies(url: str, name: str = "") -> ToolResult:
     """
-    从本地 JSON 文件加载 Cookie 到浏览器,恢复登录态。
-    Load cookies from a local JSON file into the browser.
+    根据目标 URL 自动查找本地 Cookie 文件,注入浏览器并导航到目标页面恢复登录态。
+    重要:此工具会自动完成导航,调用前不需要先调用 browser_navigate_to_url。
 
     Args:
-        name: Cookie 文件名(不含 .json 后缀)
-        url: 加载后导航到的 URL(可选)
-
-    Returns:
-        ToolResult: 加载结果
+        url: 目标 URL(必须提供,同时用于自动匹配 Cookie 文件)
+        name: Cookie 文件名(可选,不传则根据 URL 域名自动查找)
     """
     try:
         browser, tools = await get_browser_session()
 
-        cookie_file = _COOKIES_DIR / f"{name}.json"
-        if not cookie_file.exists():
-            # 列出可用文件
-            available = [f.stem for f in _COOKIES_DIR.glob("*.json")] if _COOKIES_DIR.exists() else []
-            return ToolResult(
-                title="Cookie 文件不存在",
-                output=f"未找到 .cookies/{name}.json,可用: {available}",
-                error=f"Cookie file not found: {name}.json",
-                long_term_memory=f"Cookie 文件 {name}.json 不存在"
-            )
+        if not url.startswith("http"):
+            url = f"https://{url}"
+
+        # 根据域名自动查找 Cookie 文件
+        if not name:
+            from urllib.parse import urlparse
+            domain = urlparse(url).netloc.replace("www.", "")
+            if _COOKIES_DIR.exists():
+                matches = list(_COOKIES_DIR.glob(f"{domain}*.json"))
+                if matches:
+                    cookie_file = matches[0]  # 取第一个匹配的
+                else:
+                    available = [f.stem for f in _COOKIES_DIR.glob("*.json")]
+                    return ToolResult(title="未找到 Cookie", output=f"没有匹配 {domain} 的文件,可用: {available}", error=f"无 {domain} 的 Cookie 文件")
+            else:
+                return ToolResult(title="未找到 Cookie", output=".cookies 目录不存在", error="Cookie 目录不存在")
+        else:
+            cookie_file = _COOKIES_DIR / f"{name}.json"
+            if not cookie_file.exists():
+                available = [f.stem for f in _COOKIES_DIR.glob("*.json")] if _COOKIES_DIR.exists() else []
+                return ToolResult(title="文件不存在", output=f"可用: {available}", error=f"未找到 .cookies/{name}.json")
 
         cookies = json.loads(cookie_file.read_text(encoding="utf-8"))
+
+        # 直接注入(export 和 load 使用相同的 CDP 格式,无需标准化)
         await browser._cdp_set_cookies(cookies)
 
+        # 导航到目标页面(带上刚注入的 Cookie)
         if url:
+            if not url.startswith("http"):
+                url = f"https://{url}"
             await tools.navigate(url=url, browser_session=browser)
-            await tools.wait(seconds=2, browser_session=browser)
+            await tools.wait(seconds=3, browser_session=browser)
 
         return ToolResult(
-            title="Cookie 已加载",
-            output=f"已注入 {len(cookies)} 条 Cookie(来自 {name}.json)",
-            long_term_memory=f"从 .cookies/{name}.json 加载了 {len(cookies)} 条 Cookie"
+            title="Cookie 注入并导航完成",
+            output=f"从 {cookie_file.name} 注入 {len(cookies)} 条 Cookie,已导航到 {url}",
+            long_term_memory=f"已从 {cookie_file.name} 注入 Cookie 并导航到 {url},登录态已恢复"
         )
-
     except Exception as e:
-        return ToolResult(
-            title="Cookie 加载失败",
-            output="",
-            error=f"Failed to load cookies: {str(e)}",
-            long_term_memory="加载 Cookie 失败"
-        )
+        return ToolResult(title="Cookie 加载失败", output="", error=str(e), long_term_memory="加载 Cookie 失败")
 
 
 # ============================================================

+ 1 - 1
examples/research/test.prompt

@@ -7,4 +7,4 @@ $system$
 你是最顶尖的AI助手,可以拆分并调用工具逐步解决复杂问题。
 
 $user$
-使用浏览器帮我做一下调研,打开小红书的官网,自动登录(输入手机号15035599703),点击同意协议,并点击验证,等待我手动输入(可以通过get_visual_selector_map来确定屏幕效果),然后点击登录,并随机查找一个摄影博主的最近发帖的信息
+使用浏览器,用load_cookies打开bilibili,然后搜索影视飓风

+ 30 - 0
export.py

@@ -0,0 +1,30 @@
+import asyncio
+import json
+from playwright.async_api import async_playwright
+
+async def export_cookies():
+    async with async_playwright() as p:
+        # 启动浏览器,headless=False 方便手动扫码
+        browser = await p.chromium.launch(headless=False)
+        context = await browser.new_context()
+        page = await context.new_page()
+
+        await page.goto("https://www.bilibili.com")
+        
+        print("请在浏览器中完成登录(扫码或账号密码)...")
+        
+        # 循环检查是否登录成功(根据头像或特定元素判断)
+        # 这里我们等待用户手动在控制台回车确认登录已完成
+        input("完成登录后,请按回车键继续导出 Cookie...")
+
+        # 获取所有 Cookie
+        cookies = await context.cookies()
+        
+        with open("bilibili.json", "w", encoding="utf-8") as f:
+            json.dump(cookies, f, ensure_ascii=False, indent=4)
+            
+        print("Cookie 已成功保存至 bilibili.json")
+        await browser.close()
+
+if __name__ == "__main__":
+    asyncio.run(export_cookies())

+ 33 - 0
login.py

@@ -0,0 +1,33 @@
+import asyncio
+import json
+from playwright.async_api import async_playwright
+
+async def login_with_cookies():
+    async with async_playwright() as p:
+        browser = await p.chromium.launch(headless=False) # 设为 False 以便观察效果
+        context = await browser.new_context()
+        
+        # 加载并注入 Cookie
+        try:
+            with open("bilibili.json", "r", encoding="utf-8") as f:
+                cookies = json.load(f)
+            await context.add_cookies(cookies)
+            print("Cookie 加载成功")
+        except FileNotFoundError:
+            print("找不到 Cookie 文件,请先运行导出脚本")
+            return
+
+        page = await context.new_page()
+        # 访问 B 站
+        await page.goto("https://www.bilibili.com")
+        
+        # 验证是否显示了用户名/头像(证明登录成功)
+        await page.wait_for_timeout(5000)  # 停留 5 秒观察效果
+        
+        print("当前页面标题:", await page.title())
+        # 这里可以继续你的自动化操作...
+        
+        await browser.close()
+
+if __name__ == "__main__":
+    asyncio.run(login_with_cookies())