feat(sdk): add cookies-only context.storageState helpers - #2697
Open
HaoChiBao wants to merge 5 commits into
Open
feat(sdk): add cookies-only context.storageState helpers#2697HaoChiBao wants to merge 5 commits into
HaoChiBao wants to merge 5 commits into
Conversation
Export and restore Playwright-compatible cookie storage state over CDP (including HttpOnly), with TS/Python/Go parity, docs, and tests. localStorage origins are reserved for later.
🦋 Changeset detectedLatest commit: 9d1c05a The changes in this PR will be included in the next version bump. This PR includes changesets to release 7 packages
Not sure what this means? Click here to learn what changesets are. Click here if you're a maintainer who wants to add another changeset to this PR |
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
8 issues found across 13 files
Confidence score: 2/5
- In
packages/sdk-go/browser_context.goandpackages/sdk-go/storage_state.go,SetStorageStatecan clear existing cookies before rejecting invalid input, and missinghttpOnly/securefields are treated asfalse; that combination risks silent session weakening or accidental logout on bad payloads — validate required cookie fields (includingsameSite) beforeClearCookiesand reject missing booleans explicitly. packages/sdk-python/src/stagehand/browser_context.pybreaks a round-trip flow (storage_state()output fed back intoset_storage_state()), raisingTypeErrorfor pydanticCookieentries; this is a concrete runtime regression for normal SDK usage — accept model instances (or coerce them) alongside mappings in_normalize_storage_cookie.packages/sdk-ts/src/browserContext.tswrites storage-state files with default Node permissions, so files containing HttpOnly/session cookies may be world-readable on shared systems — write with restrictive permissions (for example0o600) to reduce token exposure risk.packages/docs/v4/reference/context.mdxcurrently has a Go H2 heading that violates the enforced reference-heading regex and a Go snippet that won’t compile as pasted; this can fail docs validation and ship broken examples — align the heading with the test format and replace unusedstatewith_in the snippet.
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="packages/sdk-go/browser_context_test.go">
<violation number="1" location="packages/sdk-go/browser_context_test.go:120">
P3: The new storage-state logic has untested edge branches: `SetStorageState` returns an error when `Cookies == nil`, `normalizeStorageStateFile` returns errors for a missing `cookies` array and for invalid `sameSite` values, and `StorageState` skips the file write when `options == nil`. The added round-trip test only covers the happy path, so these newly-added error/validation branches have no coverage in any package test file (there is no `storage_state_test.go`). Add focused cases for: a file with an invalid `sameSite` (rejects), a file without a `cookies` array (rejects), and `StorageState` with `nil` options (no write, `Origins` still empty).</violation>
</file>
<file name="packages/docs/v4/reference/context.mdx">
<violation number="1" location="packages/docs/v4/reference/context.mdx:300">
P3: This sentence uses passive voice in public docs (`origins` “is reserved” / “is always”). Rewrite it in active voice with an explicit actor to match the documentation style policy.
(Based on your team's feedback about active-voice requirements in docs.) .</violation>
<violation number="2" location="packages/docs/v4/reference/context.mdx:1052">
P2: The Go `StorageState` example declares `state` but never uses it, so the snippet fails to compile when users paste it into a function. Replace `state` with `_` to keep the example buildable.</violation>
<violation number="3" location="packages/docs/v4/reference/context.mdx:1078">
P1: The Go `## SetStorageState() / SetStorageStatePath()` heading fails the enforced heading format in sdk-reference.test.ts, which parses every depth-2 heading in a language tab with the regex `^([A-Za-z_$][A-Za-z0-9_$]*)\(\)$` and throws when a heading is not an exact single method name. This throws out of readReferenceMethods() and breaks the whole docs test suite. Split it into two separate depth-2 headings, `## SetStorageState()` and `## SetStorageStatePath()`, each with its own ParamField and ResponseField, matching the two distinct Go SDK methods.</violation>
</file>
<file name="packages/sdk-go/browser_context.go">
<violation number="1" location="packages/sdk-go/browser_context.go:166">
P2: When `SetStorageState` receives an object with invalid cookie values, it clears existing cookies before the bad input is rejected downstream. Validate cookie fields (at least `SameSite`) before `ClearCookies` so invalid input fails without mutating the current cookie jar.</violation>
</file>
<file name="packages/sdk-ts/src/browserContext.ts">
<violation number="1" location="packages/sdk-ts/src/browserContext.ts:246">
P2: The exported storage-state file contains cookies including HttpOnly/session tokens, but `writeFile` uses Node's default mode (world-readable), unlike the Go SDK which deliberately writes `0o600`. Restrict permissions (e.g. `{ mode: 0o600 }`) and prefer `O_EXCL`/refuse insecure existing files where practical, mirroring the Go implementation.</violation>
</file>
<file name="packages/sdk-python/src/stagehand/browser_context.py">
<violation number="1" location="packages/sdk-python/src/stagehand/browser_context.py:275">
P2: Passing the object returned by `storage_state()` back into `set_storage_state()` raises TypeError: its `cookies` entries are pydantic `Cookie` models, but `_normalize_storage_cookie` only accepts `Mapping` objects and calls `.get(...)`/subscript on them, which pydantic `BaseModel` instances do not support. Normalize `Cookie` models too (e.g. via `Cookie.model_dump()` if it is a `Cookie`/`BaseModel`) so the documented object round-trip works like the TS/Go SDKs.</violation>
</file>
<file name="packages/sdk-go/storage_state.go">
<violation number="1" location="packages/sdk-go/storage_state.go:45">
P2: Missing `httpOnly` or `secure` in a storage-state file is silently treated as `false`, so invalid input can import weaker cookies instead of being rejected. Decode required cookie fields with presence-aware types (for example `*bool`) and return an error when they are absent before constructing `Cookie`.</violation>
</file>
Architecture diagram
sequenceDiagram
participant Client as SDK Client (TS/Python/Go)
participant BC as BrowserContext
participant CDP as CDP Cookie API
participant FS as File System (Node/Python/Go)
participant LS as LocalStorage (reserved)
Note over Client,LS: Storage State Export Flow
Client->>BC: storageState() / storage_state() / StorageState()
BC->>CDP: Get all cookies (includes HttpOnly)
CDP-->>BC: Cookie list
BC->>BC: Build StorageState {cookies, origins: []}
opt path option provided
BC->>FS: Write JSON file (0600 perms in Go)
FS-->>BC: File written
end
BC-->>Client: StorageState object
Note over BC,LS: Origins always empty (localStorage reserved)
Note over Client,LS: Storage State Import Flow
Client->>BC: setStorageState() / set_storage_state() / SetStorageState()
alt State object passed
Client->>BC: StorageState object
else File path passed
Client->>FS: Read JSON file
FS-->>BC: Raw JSON
BC->>BC: Parse and normalize (camelCase/snake_case)
end
BC->>BC: Validate cookies array present
BC->>CDP: clear_cookies() - remove existing
CDP-->>BC: Cleared
alt Cookies non-empty
BC->>BC: Map cookies to CookieParam (preserve httpOnly, secure, sameSite)
BC->>CDP: add_cookies() with params
CDP-->>BC: Cookies restored
else Cookies empty
Note over BC,CDP: Skip add_cookies, state cleared
end
BC-->>Client: Void/error (nil in Go)
Note over BC,LS: Origins entries ignored on import (localStorage not supported yet)
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
Split Go docs headings, validate cookies before clearing, require present httpOnly/secure on import, write state files as 0o600, and accept Python Cookie model round-trips.
Contributor
There was a problem hiding this comment.
All reported issues were addressed across 8 files (changes from recent commits).
Requires human review: Auto-approval blocked because this review re-detected 1 unresolved issue already reported by Cubic.
Re-trigger cubic
Reuse validateStorageStateCookies in writeStorageStateFile so export rejects invalid sameSite values the same way import does.
Contributor
Author
|
Addressed the remaining Cubic note: |
Contributor
There was a problem hiding this comment.
All reported issues were addressed across 2 files (changes from recent commits).
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
CDP cookies may omit SameSite; map empty values to Lax before validating or writing so StorageState(path) and object round-trips keep working while still rejecting invalid values.
Contributor
Author
|
Addressed the latest Cubic P2: empty |
Drop unused origins parsing, merge Go SameSite normalize/validate, and normalize empty SameSite to Lax on export/import for TS/Python parity.
Contributor
There was a problem hiding this comment.
2 issues found across 5 files (changes from recent commits).
Confidence score: 3/5
- In
packages/sdk-ts/src/browserContext.ts,readStorageStateFilecurrently rethrowsJSON.parsefailures with raw parser text, which can leak snippets of user-provided storage-state content in error surfaces and logs; map parse failures to a sanitized, fixed message while preserving only safe context (e.g., file path). - In
packages/sdk-python/src/stagehand/browser_context.py,storage_state()validates cookies before SameSite normalization, so cookies with empty/missing SameSite can still trigger model validation errors and break state export for otherwise recoverable inputs; normalize SameSite values beforeCookiemodel validation.
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="packages/sdk-ts/src/browserContext.ts">
<violation number="1" location="packages/sdk-ts/src/browserContext.ts:170">
P2: Custom agent: **Exception and error message sanitization**
The `readStorageStateFile` helper embeds the raw `JSON.parse` error message into a `TypeError` raised to the user. V8's `JSON.parse` errors can include raw input snippets (e.g. `Unexpected token 'o', "not json at"... is not valid JSON`), and storage state files contain cookies with auth/session tokens. A malformed file can therefore leak a portion of those secrets in the exception message. Replace the appended `error.message` with a generic, non-echoing description.</violation>
</file>
<file name="packages/sdk-python/src/stagehand/browser_context.py">
<violation number="1" location="packages/sdk-python/src/stagehand/browser_context.py:186">
P2: When `context.cookies` returns a cookie with missing or empty SameSite, `storage_state()` still fails before normalization because `await self.cookies()` is validated first. Normalize SameSite before `Cookie` model validation (or relax the wire model and normalize immediately after decode) so this edge case is actually recovered to `Lax`.</violation>
</file>
Tip: Review your code locally with the cubic CLI to iterate faster.
Re-trigger cubic
| function normalizeStorageStateCookieSameSite(cookies: Cookie[]): Cookie[] { | ||
| return cookies.map((cookie) => | ||
| cookie.sameSite ? cookie : { ...cookie, sameSite: "Lax" }, | ||
| ); |
Contributor
There was a problem hiding this comment.
P2: Custom agent: Exception and error message sanitization
The readStorageStateFile helper embeds the raw JSON.parse error message into a TypeError raised to the user. V8's JSON.parse errors can include raw input snippets (e.g. Unexpected token 'o', "not json at"... is not valid JSON), and storage state files contain cookies with auth/session tokens. A malformed file can therefore leak a portion of those secrets in the exception message. Replace the appended error.message with a generic, non-echoing description.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At packages/sdk-ts/src/browserContext.ts, line 170:
<comment>The `readStorageStateFile` helper embeds the raw `JSON.parse` error message into a `TypeError` raised to the user. V8's `JSON.parse` errors can include raw input snippets (e.g. `Unexpected token 'o', "not json at"... is not valid JSON`), and storage state files contain cookies with auth/session tokens. A malformed file can therefore leak a portion of those secrets in the exception message. Replace the appended `error.message` with a generic, non-echoing description.</comment>
<file context>
@@ -163,6 +163,13 @@ function cookieToParam(cookie: Cookie): CookieParam {
+function normalizeStorageStateCookieSameSite(cookies: Cookie[]): Cookie[] {
+ return cookies.map((cookie) =>
+ cookie.sameSite ? cookie : { ...cookie, sameSite: "Lax" },
+ );
+}
+
</file context>
|
|
||
| localStorage / IndexedDB are not included yet (``origins`` is always ``[]``). | ||
| """ | ||
| cookies = _normalize_storage_state_cookie_same_site(await self.cookies()) |
Contributor
There was a problem hiding this comment.
P2: When context.cookies returns a cookie with missing or empty SameSite, storage_state() still fails before normalization because await self.cookies() is validated first. Normalize SameSite before Cookie model validation (or relax the wire model and normalize immediately after decode) so this edge case is actually recovered to Lax.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At packages/sdk-python/src/stagehand/browser_context.py, line 186:
<comment>When `context.cookies` returns a cookie with missing or empty SameSite, `storage_state()` still fails before normalization because `await self.cookies()` is validated first. Normalize SameSite before `Cookie` model validation (or relax the wire model and normalize immediately after decode) so this edge case is actually recovered to `Lax`.</comment>
<file context>
@@ -183,7 +183,7 @@ async def storage_state(self, *, path: str | Path | None = None) -> StorageState
localStorage / IndexedDB are not included yet (``origins`` is always ``[]``).
"""
- cookies = await self.cookies()
+ cookies = _normalize_storage_state_cookie_same_site(await self.cookies())
state: StorageState = {"cookies": cookies, "origins": []}
if path is not None:
</file context>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
storageState/setStorageStatehelpers onBrowserContextin TypeScript, Python, and Go.originsis always empty for now (localStorage reserved).userDataDir/ Browserbase contexts already cover full-profile persistence.Why
v4 can persist a Chrome profile and CRUD cookies, but there is still no portable save/load helper for auth cookies. This is the cookies-only MVP discussed for #1250 (not a full localStorage/IndexedDB port).
Test plan
uv run pytest tests/test_browser_context.py(Python) — passed locallypnpm --filter @browserbasehq/stagehand run test:unit -- packages/sdk-ts/tests/storageState.test.tspnpm --filter @browserbasehq/stagehand exec vitest run packages/sdk-ts/tests/integration/cookies.test.ts(or CI integration)go test ./... -run StorageStateinpackages/sdk-gocontext.mdx/user-data.mdxDiscord note
Opening this for early review; will ping Discord with the PR for API-shape feedback (cookies-only vs full Playwright parity).
Summary by cubic
Adds Playwright-compatible, cookies-only storage state helpers on BrowserContext across
@browserbasehq/stagehand,@browserbasehq/stagehand-python, and@browserbasehq/stagehand-goto export/import auth cookies as portable files. Exports include HttpOnly via CDP; localStorage is not included (originsis empty). All SDKs normalize missing/empty SameSite to Lax and validate SameSite values.packages/sdk-ts): addsstorageState({ path? })andsetStorageState(stateOrPath); Node-only path I/O with clear errors; writes JSON with mode 0600; import accepts camelCase and snake_case; normalizes/validates SameSite on export/import.packages/sdk-python): addsstorage_state(path=None)andset_storage_state(state_or_path); writes pretty JSON with newline and mode 0o600; acceptsCookiemodel, mapping, or JSON path; import accepts camelCase and snake_case; normalizes/validates SameSite.packages/sdk-go): addsStorageState(ctx, *StorageStateOptions),SetStorageState(ctx, state),SetStorageStatePath(ctx, path); validates cookies before clearing; requires presenthttpOnly/secureand validsameSiteon import; normalizes empty SameSite to Lax on export/import; writes files with 0600 perms.Written for commit 9d1c05a. Summary will update on new commits.