Enable auth cookie injection

Author: rcholic

screenshot.png

@@ -32,11 +32,15 @@
     ActionTokenUsage,
     AgentActionResult,
     BBox,
+    Cookie,
     Element,
+    LocalStorageItem,
+    OriginStorage,
     ScreenshotConfig,
     Snapshot,
     SnapshotFilter,
     SnapshotOptions,
+    StorageState,
     TokenStats,
     Viewport,
     WaitResult,
@@ -54,6 +58,7 @@
     canonical_snapshot_loose,
     canonical_snapshot_strict,
     compute_snapshot_digests,
+    save_storage_state,
     sha256_digest,
 )
 from .wait import wait_for
@@ -105,6 +110,11 @@
     "SnapshotOptions",
     "SnapshotFilter",
     "ScreenshotConfig",
+    # Storage State Models (Auth Injection)
+    "StorageState",
+    "Cookie",
+    "LocalStorageItem",
+    "OriginStorage",
     # Tracing (v0.12.0+)
     "Tracer",
     "TraceSink",
@@ -118,6 +128,7 @@
     "canonical_snapshot_loose",
     "compute_snapshot_digests",
     "sha256_digest",
+    "save_storage_state",
     # Formatting (v0.12.0+)
     "format_snapshot_for_llm",
     # Agent Config (v0.12.0+)

sentience/__init__.py

@@ -11,7 +11,7 @@
 
 from playwright.sync_api import BrowserContext, Page, Playwright, sync_playwright
 
-from sentience.models import ProxyConfig
+from sentience.models import ProxyConfig, StorageState
 
 # Import stealth for bot evasion (optional - graceful fallback if not available)
 try:
@@ -31,6 +31,8 @@ def __init__(
         api_url: str | None = None,
         headless: bool | None = None,
         proxy: str | None = None,
+        user_data_dir: str | None = None,
+        storage_state: str | Path | StorageState | dict | None = None,
     ):
         """
         Initialize Sentience browser
@@ -46,6 +48,15 @@ def __init__(
             proxy: Optional proxy server URL (e.g., 'http://user:pass@proxy.example.com:8080')
                    Supports HTTP, HTTPS, and SOCKS5 proxies
                    Falls back to SENTIENCE_PROXY environment variable if not provided
+            user_data_dir: Optional path to user data directory for persistent sessions.
+                          If None, uses temporary directory (session not persisted).
+                          If provided, cookies and localStorage persist across browser restarts.
+            storage_state: Optional storage state to inject (cookies + localStorage).
+                          Can be:
+                          - Path to JSON file (str or Path)
+                          - StorageState object
+                          - Dictionary with 'cookies' and/or 'origins' keys
+                          If provided, browser starts with pre-injected authentication.
         """
         self.api_key = api_key
         # Only set api_url if api_key is provided, otherwise None (free tier)
@@ -65,6 +76,10 @@ def __init__(
         # Support proxy from argument or environment variable
         self.proxy = proxy or os.environ.get("SENTIENCE_PROXY")
 
+        # Auth injection support
+        self.user_data_dir = user_data_dir
+        self.storage_state = storage_state
+
         self.playwright: Playwright | None = None
         self.context: BrowserContext | None = None
         self.page: Page | None = None
@@ -170,9 +185,16 @@ def start(self) -> None:
         # Parse proxy configuration if provided
         proxy_config = self._parse_proxy(self.proxy) if self.proxy else None
 
+        # Handle User Data Directory (Persistence)
+        if self.user_data_dir:
+            user_data_dir = str(self.user_data_dir)
+            Path(user_data_dir).mkdir(parents=True, exist_ok=True)
+        else:
+            user_data_dir = ""  # Ephemeral temp dir (existing behavior)
+
         # Build launch_persistent_context parameters
         launch_params = {
-            "user_data_dir": "",  # Ephemeral temp dir
+            "user_data_dir": user_data_dir,
             "headless": False,  # IMPORTANT: See note above
             "args": args,
             "viewport": {"width": 1280, "height": 800},
@@ -194,6 +216,10 @@ def start(self) -> None:
 
         self.page = self.context.pages[0] if self.context.pages else self.context.new_page()
 
+        # Inject storage state if provided (must be after context creation)
+        if self.storage_state:
+            self._inject_storage_state(self.storage_state)
+
         # Apply stealth if available
         if STEALTH_AVAILABLE:
             stealth_sync(self.page)
@@ -233,6 +259,92 @@ def goto(self, url: str) -> None:
                 f"5. Diagnostic info: {diag}"
             )
 
+    def _inject_storage_state(
+        self, storage_state: str | Path | StorageState | dict
+    ) -> None:  # noqa: C901
+        """
+        Inject storage state (cookies + localStorage) into browser context.
+
+        Args:
+            storage_state: Path to JSON file, StorageState object, or dict containing storage state
+        """
+        import json
+
+        # Load storage state
+        if isinstance(storage_state, (str, Path)):
+            # Load from file
+            with open(storage_state, encoding="utf-8") as f:
+                state_dict = json.load(f)
+            state = StorageState.from_dict(state_dict)
+        elif isinstance(storage_state, StorageState):
+            # Already a StorageState object
+            state = storage_state
+        elif isinstance(storage_state, dict):
+            # Dictionary format
+            state = StorageState.from_dict(storage_state)
+        else:
+            raise ValueError(
+                f"Invalid storage_state type: {type(storage_state)}. "
+                "Expected str, Path, StorageState, or dict."
+            )
+
+        # Inject cookies (works globally)
+        if state.cookies:
+            # Convert to Playwright cookie format
+            playwright_cookies = []
+            for cookie in state.cookies:
+                cookie_dict = cookie.model_dump()
+                # Playwright expects lowercase keys for some fields
+                playwright_cookie = {
+                    "name": cookie_dict["name"],
+                    "value": cookie_dict["value"],
+                    "domain": cookie_dict["domain"],
+                    "path": cookie_dict["path"],
+                }
+                if cookie_dict.get("expires"):
+                    playwright_cookie["expires"] = cookie_dict["expires"]
+                if cookie_dict.get("httpOnly"):
+                    playwright_cookie["httpOnly"] = cookie_dict["httpOnly"]
+                if cookie_dict.get("secure"):
+                    playwright_cookie["secure"] = cookie_dict["secure"]
+                if cookie_dict.get("sameSite"):
+                    playwright_cookie["sameSite"] = cookie_dict["sameSite"]
+                playwright_cookies.append(playwright_cookie)
+
+            self.context.add_cookies(playwright_cookies)
+            print(f"✅ [Sentience] Injected {len(state.cookies)} cookie(s)")
+
+        # Inject LocalStorage (requires navigation to each domain)
+        if state.origins:
+            for origin_data in state.origins:
+                origin = origin_data.origin
+                if not origin:
+                    continue
+
+                # Navigate to origin to set localStorage
+                try:
+                    self.page.goto(origin, wait_until="domcontentloaded", timeout=10000)
+
+                    # Inject localStorage
+                    if origin_data.localStorage:
+                        # Convert to dict format for JavaScript
+                        localStorage_dict = {
+                            item.name: item.value for item in origin_data.localStorage
+                        }
+                        self.page.evaluate(
+                            """(localStorage_data) => {
+                                for (const [key, value] of Object.entries(localStorage_data)) {
+                                    localStorage.setItem(key, value);
+                                }
+                            }""",
+                            localStorage_dict,
+                        )
+                        print(
+                            f"✅ [Sentience] Injected {len(origin_data.localStorage)} localStorage item(s) for {origin}"
+                        )
+                except Exception as e:
+                    print(f"⚠️  [Sentience] Failed to inject localStorage for {origin}: {e}")
+
     def _wait_for_extension(self, timeout_sec: float = 5.0) -> bool:
         """Poll for window.sentience to be available"""
         start_time = time.time()

sentience/browser.py

@@ -216,3 +216,124 @@ def to_playwright_dict(self) -> dict:
             config["username"] = self.username
             config["password"] = self.password
         return config
+
+
+# ========== Storage State Models (Auth Injection) ==========
+
+
+class Cookie(BaseModel):
+    """
+    Cookie definition for storage state injection.
+
+    Matches Playwright's cookie format for storage_state.
+    """
+
+    name: str = Field(..., description="Cookie name")
+    value: str = Field(..., description="Cookie value")
+    domain: str = Field(..., description="Cookie domain (e.g., '.example.com')")
+    path: str = Field(default="/", description="Cookie path")
+    expires: float | None = Field(None, description="Expiration timestamp (Unix epoch)")
+    httpOnly: bool = Field(default=False, description="HTTP-only flag")
+    secure: bool = Field(default=False, description="Secure (HTTPS-only) flag")
+    sameSite: Literal["Strict", "Lax", "None"] = Field(
+        default="Lax", description="SameSite attribute"
+    )
+
+
+class LocalStorageItem(BaseModel):
+    """
+    LocalStorage item for a specific origin.
+
+    Playwright stores localStorage as an array of {name, value} objects.
+    """
+
+    name: str = Field(..., description="LocalStorage key")
+    value: str = Field(..., description="LocalStorage value")
+
+
+class OriginStorage(BaseModel):
+    """
+    Storage state for a specific origin (localStorage).
+
+    Represents localStorage data for a single domain.
+    """
+
+    origin: str = Field(..., description="Origin URL (e.g., 'https://example.com')")
+    localStorage: list[LocalStorageItem] = Field(
+        default_factory=list, description="LocalStorage items for this origin"
+    )
+
+
+class StorageState(BaseModel):
+    """
+    Complete browser storage state (cookies + localStorage).
+
+    This is the format used by Playwright's storage_state() method.
+    Can be saved to/loaded from JSON files for session injection.
+    """
+
+    cookies: list[Cookie] = Field(
+        default_factory=list, description="Cookies to inject (global scope)"
+    )
+    origins: list[OriginStorage] = Field(
+        default_factory=list, description="LocalStorage data per origin"
+    )
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "StorageState":
+        """
+        Create StorageState from dictionary (e.g., loaded from JSON).
+
+        Args:
+            data: Dictionary with 'cookies' and/or 'origins' keys
+
+        Returns:
+            StorageState instance
+        """
+        cookies = [
+            Cookie(**cookie) if isinstance(cookie, dict) else cookie
+            for cookie in data.get("cookies", [])
+        ]
+        origins = []
+        for origin_data in data.get("origins", []):
+            if isinstance(origin_data, dict):
+                # Handle localStorage as array of {name, value} or as dict
+                localStorage_data = origin_data.get("localStorage", [])
+                if isinstance(localStorage_data, dict):
+                    # Convert dict to list of LocalStorageItem
+                    localStorage_items = [
+                        LocalStorageItem(name=k, value=v) for k, v in localStorage_data.items()
+                    ]
+                else:
+                    # Already a list
+                    localStorage_items = [
+                        LocalStorageItem(**item) if isinstance(item, dict) else item
+                        for item in localStorage_data
+                    ]
+                origins.append(
+                    OriginStorage(
+                        origin=origin_data.get("origin", ""),
+                        localStorage=localStorage_items,
+                    )
+                )
+            else:
+                origins.append(origin_data)
+        return cls(cookies=cookies, origins=origins)
+
+    def to_playwright_dict(self) -> dict:
+        """
+        Convert to Playwright-compatible dictionary format.
+
+        Returns:
+            Dictionary compatible with Playwright's storage_state parameter
+        """
+        return {
+            "cookies": [cookie.model_dump() for cookie in self.cookies],
+            "origins": [
+                {
+                    "origin": origin.origin,
+                    "localStorage": [item.model_dump() for item in origin.localStorage],
+                }
+                for origin in self.origins
+            ],
+        }

sentience/models.py

@@ -11,7 +11,10 @@
 import json
 import re
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+from pathlib import Path
+from typing import Any
+
+from playwright.sync_api import BrowserContext
 
 
 @dataclass
@@ -255,3 +258,39 @@ def compute_snapshot_digests(elements: list[dict[str, Any]]) -> dict[str, str]:
         "strict": sha256_digest(canonical_strict),
         "loose": sha256_digest(canonical_loose),
     }
+
+
+def save_storage_state(context: BrowserContext, file_path: str | Path) -> None:
+    """
+    Save current browser storage state (cookies + localStorage) to a file.
+
+    This is useful for capturing a logged-in session to reuse later.
+
+    Args:
+        context: Playwright BrowserContext
+        file_path: Path to save the storage state JSON file
+
+    Example:
+        ```python
+        from sentience import SentienceBrowser, save_storage_state
+
+        browser = SentienceBrowser()
+        browser.start()
+
+        # User logs in manually or via agent
+        browser.goto("https://example.com")
+        # ... login happens ...
+
+        # Save session for later
+        save_storage_state(browser.context, "auth.json")
+        ```
+
+    Raises:
+        IOError: If file cannot be written
+    """
+    storage_state = context.storage_state()
+    file_path_obj = Path(file_path)
+    file_path_obj.parent.mkdir(parents=True, exist_ok=True)
+    with open(file_path_obj, "w") as f:
+        json.dump(storage_state, f, indent=2)
+    print(f"✅ [Sentience] Saved storage state to {file_path_obj}")