fix domain based sensitive_data

This commit is contained in:
Anirudha619
2025-06-06 00:04:33 +05:30
parent 8a1d3a6245
commit 5033d643cc
2 changed files with 27 additions and 19 deletions

View File

@@ -17,7 +17,7 @@ from browser_use.agent.message_manager.views import MessageMetadata
from browser_use.agent.prompts import AgentMessagePrompt
from browser_use.agent.views import ActionResult, AgentOutput, AgentStepInfo, MessageManagerState
from browser_use.browser.views import BrowserStateSummary
from browser_use.utils import time_execution_sync
from browser_use.utils import match_url_with_domain_pattern, time_execution_sync
logger = logging.getLogger(__name__)
@@ -224,24 +224,6 @@ class MessageManager:
)
self._add_message_with_tokens(task_message, message_type='init')
if self.settings.sensitive_data:
sensitive_data = self.settings.sensitive_data
# Check if any value is a dict (domain-based sensitive data)
if any(isinstance(v, dict) for v in sensitive_data.values()):
# Domain-based: collect all keys from all value dicts
placeholders = set()
for v in sensitive_data.values():
if isinstance(v, dict):
placeholders.update(v.keys())
print('placeholders ', placeholders)
info = f'Here are placeholders for sensitive data: {list(placeholders)}'
else:
# Old format: just use the top-level keys
info = f'Here are placeholders for sensitive data: {list(sensitive_data.keys())}'
info += '\nTo use them, write <secret>the placeholder name</secret>'
info_message = HumanMessage(content=info)
self._add_message_with_tokens(info_message, message_type='init')
placeholder_message = HumanMessage(content='Example output:')
self._add_message_with_tokens(placeholder_message, message_type='init')
@@ -294,6 +276,29 @@ class MessageManager:
self._add_message_with_tokens(msg)
self.task = new_task
def add_sensitive_data(self, current_page_url) -> None:
sensitive_data = self.settings.sensitive_data
if not sensitive_data:
return
# Collect placeholders for sensitive data
placeholders = set()
for key, value in sensitive_data.items():
if isinstance(value, dict):
# New format: {domain: {key: value}}
if match_url_with_domain_pattern(current_page_url, key, True):
placeholders.update(value.keys())
else:
# Old format: {key: value}
placeholders.add(key)
if placeholders:
print('placeholders ', placeholders)
info = f'Here are placeholders for sensitive data: {list(placeholders)}'
info += '\nTo use them, write <secret>the placeholder name</secret>'
info_message = HumanMessage(content=info)
self._add_message_with_tokens(info_message, message_type='init')
@time_execution_sync('--add_state_message')
def add_state_message(
self,

View File

@@ -800,6 +800,9 @@ class Agent(Generic[Context]):
# Get page-specific filtered actions
page_filtered_actions = self.controller.registry.get_prompt_description(current_page)
if self.sensitive_data:
self._message_manager.add_sensitive_data(current_page.url)
# If there are page-specific actions, add them as a special message for this step only
if page_filtered_actions:
page_action_message = f'For this page, these additional actions are available:\n{page_filtered_actions}'