Skip to content

Commit

Permalink
Resume session (build off #172) (#174)
Browse files Browse the repository at this point in the history 
* browserbaseResumeSessionID to connect to an running bb session (#172)

* enforce project ID

* api key and project id as params

---------

Co-authored-by: Marko Kraemer <[email protected]>
  • Loading branch information
@kamath @markokraemer
kamath and markokraemer authored Nov 10, 2024
1 parent 4ef6eb4 commit 25f9214
Show file tree
Hide file tree
Showing 2 changed files with 98 additions and 32 deletions.
129 changes: 97 additions & 32 deletions lib/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ import { modelsWithVision } from "./llm/LLMClient";
require("dotenv").config({ path: ".env" });

async function getBrowser(
apiKey: string | undefined,
projectId: string | undefined,
env: "LOCAL" | "BROWSERBASE" = "LOCAL",
headless: boolean = false,
logger: (message: {
Expand All @@ -20,55 +22,103 @@ async function getBrowser(
level?: 0 | 1 | 2;
}) => void,
browserbaseSessionCreateParams?: Browserbase.Sessions.SessionCreateParams,
browserbaseResumeSessionID?: string,
) {
if (env === "BROWSERBASE" && !process.env.BROWSERBASE_API_KEY) {
logger({
category: "Init",
message:
"BROWSERBASE_API_KEY is required to use BROWSERBASE env. Defaulting to LOCAL.",
level: 0,
});
env = "LOCAL";
}

if (env === "BROWSERBASE" && !process.env.BROWSERBASE_PROJECT_ID) {
logger({
category: "Init",
message:
"BROWSERBASE_PROJECT_ID is required to use BROWSERBASE env. Defaulting to LOCAL.",
level: 0,
});
env = "LOCAL";
if (env === "BROWSERBASE") {
if (!apiKey) {
logger({
category: "Init",
message:
"BROWSERBASE_API_KEY is required to use BROWSERBASE env. Defaulting to LOCAL.",
level: 0,
});
env = "LOCAL";
}
if (!projectId) {
logger({
category: "Init",
message:
"BROWSERBASE_PROJECT_ID is required for some Browserbase features that may not work without it.",
level: 1,
});
}
}

if (env === "BROWSERBASE") {
if (!apiKey) {
throw new Error("BROWSERBASE_API_KEY is required.");
}

let debugUrl: string | undefined = undefined;
let sessionUrl: string | undefined = undefined;
let sessionId: string;
let connectUrl: string;

logger({
category: "Init",
message: "Connecting you to Browserbase...",
level: 0,
});
const browserbase = new Browserbase({
apiKey: process.env.BROWSERBASE_API_KEY,
});
const session = await browserbase.sessions.create({
projectId: process.env.BROWSERBASE_PROJECT_ID,
...browserbaseSessionCreateParams,
apiKey,
});

const sessionId = session.id;
const connectUrl = session.connectUrl;
const browser = await chromium.connectOverCDP(connectUrl);
if (browserbaseResumeSessionID) {
// Validate the session status
try {
const sessionStatus = await browserbase.sessions.retrieve(
browserbaseResumeSessionID,
);

if (sessionStatus.status !== "RUNNING") {
throw new Error(
`Session ${browserbaseResumeSessionID} is not running (status: ${sessionStatus.status})`,
);
}

sessionId = browserbaseResumeSessionID;
connectUrl = `wss://connect.browserbase.com?apiKey=${apiKey}&sessionId=${sessionId}`;

logger({
category: "Init",
message: "Resuming existing Browserbase session...",
level: 0,
});
} catch (error) {
logger({
category: "Init",
message: `Failed to resume session ${browserbaseResumeSessionID}: ${error.message}`,
level: 0,
});
throw error;
}
} else {
// Create new session (existing code)
logger({
category: "Init",
message: "Creating new Browserbase session...",
level: 0,
});

if (!projectId) {
throw new Error(
"BROWSERBASE_PROJECT_ID is required for new Browserbase sessions.",
);
}

const session = await browserbase.sessions.create({
projectId,
...browserbaseSessionCreateParams,
});

sessionId = session.id;
connectUrl = session.connectUrl;
}

const browser = await chromium.connectOverCDP(connectUrl);
const { debuggerUrl } = await browserbase.sessions.debug(sessionId);

debugUrl = debuggerUrl;
sessionUrl = `https://www.browserbase.com/sessions/${sessionId}`;

logger({
category: "Init",
message: `Browserbase session started.\n\nSession Url: ${sessionUrl}\n\nLive debug accessible here: ${debugUrl}.`,
message: `Browserbase session ${browserbaseResumeSessionID ? "resumed" : "started"}.\n\nSession Url: ${sessionUrl}\n\nLive debug accessible here: ${debugUrl}.`,
level: 0,
});

Expand Down Expand Up @@ -182,6 +232,8 @@ export class Stagehand {
public page: Page;
public context: BrowserContext;
private env: "LOCAL" | "BROWSERBASE";
private apiKey: string | undefined;
private projectId: string | undefined;
private verbose: 0 | 1 | 2;
private debugDom: boolean;
private defaultModelName: AvailableModel;
Expand All @@ -194,10 +246,13 @@ export class Stagehand {
private domSettleTimeoutMs: number;
private browserBaseSessionCreateParams?: Browserbase.Sessions.SessionCreateParams;
private enableCaching: boolean;
private browserbaseResumeSessionID?: string;

constructor(
{
env,
apiKey,
projectId,
verbose,
debugDom,
llmProvider,
Expand All @@ -206,8 +261,11 @@ export class Stagehand {
browserBaseSessionCreateParams,
domSettleTimeoutMs,
enableCaching,
browserbaseResumeSessionID,
}: {
env: "LOCAL" | "BROWSERBASE";
apiKey?: string;
projectId?: string;
verbose?: 0 | 1 | 2;
debugDom?: boolean;
llmProvider?: LLMProvider;
Expand All @@ -220,6 +278,7 @@ export class Stagehand {
domSettleTimeoutMs?: number;
browserBaseSessionCreateParams?: Browserbase.Sessions.SessionCreateParams;
enableCaching?: boolean;
browserbaseResumeSessionID?: string;
} = {
env: "BROWSERBASE",
},
Expand All @@ -231,13 +290,16 @@ export class Stagehand {
llmProvider || new LLMProvider(this.logger, this.enableCaching);
this.env = env;
this.observations = {};
this.apiKey = apiKey;
this.projectId = projectId;
this.actions = {};
this.verbose = verbose ?? 0;
this.debugDom = debugDom ?? false;
this.defaultModelName = "gpt-4o";
this.domSettleTimeoutMs = domSettleTimeoutMs ?? 60_000;
this.headless = headless ?? false;
this.browserBaseSessionCreateParams = browserBaseSessionCreateParams;
this.browserbaseResumeSessionID = browserbaseResumeSessionID;
}

async init({
Expand All @@ -247,10 +309,13 @@ export class Stagehand {
sessionUrl: string;
}> {
const { context, debugUrl, sessionUrl } = await getBrowser(
this.apiKey,
this.projectId,
this.env,
this.headless,
this.logger,
this.browserBaseSessionCreateParams,
this.browserbaseResumeSessionID,
).catch((e) => {
console.error("Error in init:", e);
return { context: undefined, debugUrl: undefined, sessionUrl: undefined };
Expand Down
1 change: 1 addition & 0 deletions package-lock.json
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 25f9214

Please sign in to comment.