mirror of
https://github.com/browser-use/browser-use
synced 2026-05-06 17:52:15 +02:00
105 lines
3.6 KiB
Python
105 lines
3.6 KiB
Python
"""
|
||
Custom HTTP Headers via a custom Watchdog.
|
||
|
||
Creates a custom watchdog that listens to TabCreatedEvent and injects
|
||
custom HTTP headers into every new tab using Network.setExtraHTTPHeaders.
|
||
|
||
Note: The CDP EventRegistry only supports one handler per event method,
|
||
so registering directly on Target.attachedToTarget would replace the
|
||
internal SessionManager handler. Using the browser-use event system
|
||
(TabCreatedEvent) avoids this and fires after the target is fully set up.
|
||
|
||
Note: Network.setExtraHTTPHeaders is a full replacement (not additive).
|
||
|
||
Verified by navigating to https://httpbin.org/headers in a new tab.
|
||
"""
|
||
|
||
import asyncio
|
||
import os
|
||
import sys
|
||
from typing import ClassVar
|
||
|
||
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
|
||
|
||
from bubus import BaseEvent
|
||
from dotenv import load_dotenv
|
||
|
||
load_dotenv()
|
||
|
||
from browser_use import Agent, Browser, ChatBrowserUse
|
||
from browser_use.browser.events import AgentFocusChangedEvent, TabCreatedEvent
|
||
from browser_use.browser.watchdog_base import BaseWatchdog
|
||
|
||
CUSTOM_HEADERS = {
|
||
'X-Custom-Auth': 'Bearer my-secret-token',
|
||
'X-Request-Source': 'browser-use-agent',
|
||
'X-Trace-Id': 'example-trace-12345',
|
||
}
|
||
|
||
|
||
class CustomHeadersWatchdog(BaseWatchdog):
|
||
"""Injects custom HTTP headers on every new tab and focus change.
|
||
|
||
Listens to both TabCreatedEvent (new tabs) and AgentFocusChangedEvent
|
||
(tab switches) because headers are bound to a CDP session, and sessions
|
||
can be recreated on cross-origin navigations or tab switches.
|
||
"""
|
||
|
||
LISTENS_TO: ClassVar[list[type[BaseEvent]]] = [TabCreatedEvent, AgentFocusChangedEvent]
|
||
EMITS: ClassVar[list[type[BaseEvent]]] = []
|
||
|
||
async def on_TabCreatedEvent(self, event: TabCreatedEvent) -> None:
|
||
"""Set extra headers when a new tab is created."""
|
||
try:
|
||
await self.browser_session.set_extra_headers(CUSTOM_HEADERS, target_id=event.target_id)
|
||
except Exception as e:
|
||
self.logger.debug(f'Could not set headers on {event.target_id[:8]}: {e}')
|
||
|
||
async def on_AgentFocusChangedEvent(self, event: AgentFocusChangedEvent) -> None:
|
||
"""Re-apply headers when the agent switches to a different tab."""
|
||
try:
|
||
await self.browser_session.set_extra_headers(CUSTOM_HEADERS, target_id=event.target_id)
|
||
except Exception as e:
|
||
self.logger.debug(f'Could not set headers on {event.target_id[:8]}: {e}')
|
||
|
||
|
||
async def main():
|
||
browser = Browser(headless=False)
|
||
|
||
# Start the browser so watchdogs are initialized
|
||
await browser.start()
|
||
|
||
# Attach our custom watchdog to the browser session
|
||
CustomHeadersWatchdog.model_rebuild()
|
||
headers_watchdog = CustomHeadersWatchdog(event_bus=browser.event_bus, browser_session=browser)
|
||
headers_watchdog.attach_to_session()
|
||
|
||
# The watchdog only fires for tabs created AFTER registration.
|
||
# To apply headers to an already-existing tab, call set_extra_headers():
|
||
#
|
||
# await browser.set_extra_headers(CUSTOM_HEADERS)
|
||
# await browser.set_extra_headers(CUSTOM_HEADERS, target_id=some_target_id)
|
||
#
|
||
# Keep in mind that setExtraHTTPHeaders is a full replacement – each
|
||
# call overwrites all previously set extra headers on that target.
|
||
|
||
# Run the agent – open httpbin.org/headers in a new tab so the
|
||
# watchdog fires and injects the custom headers.
|
||
agent = Agent(
|
||
task=(
|
||
'Open https://httpbin.org/headers in two different tabs and extract the full JSON response. '
|
||
'Look for the custom headers X-Custom-Auth, X-Request-Source, and X-Trace-Id in the output and compare the results.'
|
||
),
|
||
llm=ChatBrowserUse(model='bu-2-0'),
|
||
browser=browser,
|
||
)
|
||
|
||
result = await agent.run()
|
||
print(result.final_result())
|
||
|
||
await browser.kill()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
asyncio.run(main())
|