fix the session close default session creation behavior (#82)

Author: Kylejeong2

browserbase/src/context.ts

@@ -3,6 +3,7 @@ import type { BrowserSession } from "./sessionManager.js";
 import {
   getSession,
   defaultSessionId,
+  getSessionReadOnly,
 } from "./sessionManager.js";
 import type { Tool, ToolResult } from "./tools/tool.js";
 import type { Config } from "../config.js";
@@ -195,6 +196,34 @@ export class Context {
     return session.browser;
   }
 
+  /**
+   * Get the active browser without triggering session creation.
+   * This is a read-only operation used when we need to check for an existing browser
+   * without side effects (e.g., during close operations).
+   * @returns The browser if it exists and is connected, null otherwise
+   */
+  public getActiveBrowserReadOnly(): BrowserSession["browser"] | null {
+    const session = getSessionReadOnly(this.currentSessionId);
+    if (!session || !session.browser || !session.browser.isConnected()) {
+      return null;
+    }
+    return session.browser;
+  }
+
+  /**
+   * Get the active page without triggering session creation.
+   * This is a read-only operation used when we need to check for an existing page
+   * without side effects.
+   * @returns The page if it exists and is not closed, null otherwise
+   */
+  public getActivePageReadOnly(): BrowserSession["page"] | null {
+    const session = getSessionReadOnly(this.currentSessionId);
+    if (!session || !session.page || session.page.isClosed()) {
+      return null;
+    }
+    return session.page;
+  }
+
   public async waitForTimeout(timeoutMillis: number): Promise<void> {
     return new Promise((resolve) => setTimeout(resolve, timeoutMillis));
   }
@@ -507,7 +536,8 @@ export class Context {
 
           // 2. Prepare and add additional textual information (URL, Title, Snapshot)
           const additionalInfoParts: string[] = [];
-          const currentPage = await this.getActivePage();
+          // Use read-only version to avoid creating sessions after close
+          const currentPage = this.getActivePageReadOnly();
 
           if (currentPage) {
             try {

browserbase/src/sessionManager.ts

@@ -324,6 +324,64 @@ export async function getSession(
   return sessionObj;
 }
 
+/**
+ * Get a session by ID without creating new sessions.
+ * This is a read-only operation that never triggers session creation.
+ * Used for operations like closing sessions where we don't want side effects.
+ * @param sessionId The session ID to retrieve
+ * @returns The session if it exists and is valid, null otherwise
+ */
+export function getSessionReadOnly(sessionId: string): BrowserSession | null {
+  // Check if it's the default session
+  if (sessionId === defaultSessionId && defaultBrowserSession) {
+    // Only return if it's actually connected and valid
+    if (defaultBrowserSession.browser.isConnected() && !defaultBrowserSession.page.isClosed()) {
+      return defaultBrowserSession;
+    }
+    return null;
+  }
+
+  // For non-default sessions, check the browsers map
+  const sessionObj = browsers.get(sessionId);
+  if (!sessionObj) {
+    return null;
+  }
+
+  // Validate the session is still active
+  if (!sessionObj.browser.isConnected() || sessionObj.page.isClosed()) {
+    return null;
+  }
+
+  return sessionObj;
+}
+
+/**
+ * Clean up a session by removing it from tracking.
+ * This is called after a browser is closed to ensure proper cleanup.
+ * @param sessionId The session ID to clean up
+ */
+export function cleanupSession(sessionId: string): void {
+  process.stderr.write(
+    `[SessionManager] Cleaning up session: ${sessionId}\n`
+  );
+  
+  // Remove from browsers map
+  browsers.delete(sessionId);
+  
+  // Clear default session reference if this was the default
+  if (sessionId === defaultSessionId && defaultBrowserSession) {
+    defaultBrowserSession = null;
+  }
+  
+  // Reset active session to default if this was the active one
+  if (activeSessionId === sessionId) {
+    process.stderr.write(
+      `[SessionManager] Cleaned up active session ${sessionId}, resetting to default.\n`
+    );
+    setActiveSessionId(defaultSessionId);
+  }
+}
+
 // Function to close all managed browser sessions gracefully
 export async function closeAllSessions(): Promise<void> {
   process.stderr.write(`[SessionManager] Closing all sessions...\n`);

browserbase/src/tools/common.ts

@@ -28,15 +28,15 @@ async function handleWait(context: Context, params: WaitInput): Promise<ToolResu
 
 // Define tool using handle
 const waitTool: Tool<typeof WaitInputSchema> = {
-    capability: 'core', // Add capability
+    capability: 'core', 
     schema: waitSchema,
     handle: handleWait,
 };
 
 
 // --- Tool: Close ---
 const CloseInputSchema = z.object({
-    random_string: z.string().optional().describe("Dummy parameter") // Keep schema if needed
+    random_string: z.string().optional().describe("Dummy parameter") 
 });
 type CloseInput = z.infer<typeof CloseInputSchema>;
 

browserbase/src/tools/getText.ts

@@ -6,7 +6,6 @@ import type { ToolActionResult } from '../context.js';
 // --- Tool: Get Text ---
 const GetTextInputSchema = z.object({
     selector: z.string().optional().describe("Optional CSS selector to get text from. If omitted, gets text from the whole body."),
-    sessionId: z.string().optional(), // Keep for schema consistency if needed elsewhere
 });
 type GetTextInput = z.infer<typeof GetTextInputSchema>;
 

browserbase/src/tools/session.ts

@@ -8,10 +8,10 @@ import {
   createNewBrowserSession,
   defaultSessionId,
   ensureDefaultSessionInternal,
+  cleanupSession,
   type BrowserSession,
 } from "../sessionManager.js";
 
-
 // --- Tool: Create Session ---
 const CreateSessionInputSchema = z.object({
   // Keep sessionId optional, but clarify its role
@@ -43,7 +43,8 @@ async function handleCreateSession(
       let targetSessionId: string;
 
       if (params.sessionId) {
-        targetSessionId = params.sessionId;
+        const projectId = config.browserbaseProjectId || '';
+        targetSessionId = `${params.sessionId}_${projectId}`;
         process.stderr.write(
           `[tool.createSession] Attempting to create/assign session with specified ID: ${targetSessionId}`
         );
@@ -140,10 +141,10 @@ async function handleCloseSession(
     let browserClosedSuccessfully = false;
     let browserCloseErrorMessage = "";
 
-    // Step 1: Attempt to get the active browser instance
+    // Step 1: Attempt to get the active browser instance WITHOUT creating a new one
     try {
-      // This call might be associated with 'previousSessionId' (which could be default or specific)
-      browser = await context.getActiveBrowser();
+      // Use read-only version to avoid creating new sessions
+      browser = context.getActiveBrowserReadOnly();
     } catch (error: any) {
       process.stderr.write(
         `[tool.closeSession] Error retrieving active browser (session ID was ${previousSessionId || 'default/unknown'}): ${error.message || String(error)}`
@@ -164,6 +165,9 @@ async function handleCloseSession(
           `[tool.closeSession] Browser connection for session (was ${previousSessionId}) closed.`
         );
 
+        // Clean up the session from tracking
+        cleanupSession(previousSessionId);
+
         process.stderr.write(
           `[tool.closeSession] View session replay at https://www.browserbase.com/sessions/${previousSessionId}`
         );
@@ -204,7 +208,7 @@ async function handleCloseSession(
     if (browserClosedSuccessfully) { // Browser was present and closed
       let successMessage = `Browserbase session (associated with context ID ${previousSessionId || 'default'}) closed successfully. Context reset to default.`;
       if (previousSessionId && previousSessionId !== defaultSessionId) {
-        successMessage += ` If this was a uniquely named session (${previousSessionId}), view replay (if available) at https://browserbase.com/sessions/${previousSessionId}`;
+        successMessage += ` If this was a uniquely named session (${previousSessionId}), view replay (if available) at https://browserbase.com/sessions`;
       }
       return { content: [{ type: "text", text: successMessage }] };
     }

browserbase/src/tools/snapshot.ts

@@ -3,10 +3,10 @@ import type {
   TextContent,
   ImageContent,
 } from "@modelcontextprotocol/sdk/types.js";
+import type { Locator, PageScreenshotOptions } from "playwright-core";
 
 import { defineTool, type ToolResult,  } from "./tool.js";
 import type { Context, ToolActionResult } from "../context.js"; 
-import type { Page, Locator } from "playwright-core"; 
 import { PageSnapshot } from "../pageSnapshot.js"; 
 import { outputFile } from "../config.js"; 
 
@@ -358,27 +358,23 @@ const selectOption = defineTool<typeof selectOptionSchema>({
 });
 
 // --- Tool: Screenshot (Adapted Handle, Example Action) ---
-const screenshotSchema = z
-  .object({
-    raw: z
-      .boolean()
-      .optional()
-      .describe(
-        "Whether to return without compression (PNG). Default is false (JPEG)."
-      ),
-    element: z
-      .string()
-      .optional()
-      .describe("Human-readable element description."),
-    ref: z
-      .string()
-      .optional()
-      .describe("Exact target element reference from the page snapshot."),
-  })
-  .refine((data) => !!data.element === !!data.ref, {
-    message: "Both element and ref must be provided or neither.",
-    path: ["ref", "element"],
-  });
+const screenshotSchema = z.object({
+  raw: z
+    .boolean()
+    .optional()
+    .describe(
+      "Whether to return without compression (PNG). Default is false (JPEG)."
+    ),
+  element: z
+    .string()
+    .optional()
+    .describe("Human-readable element description."),
+  ref: z
+    .string()
+    .optional()
+    .describe("Exact target element reference from the page snapshot.")
+});
+
 type ScreenshotInput = z.infer<typeof screenshotSchema>;
 
 const screenshot = defineTool<typeof screenshotSchema>({
@@ -392,6 +388,10 @@ const screenshot = defineTool<typeof screenshotSchema>({
     context: Context,
     params: ScreenshotInput
   ): Promise<ToolResult> => {
+    if (!!params.element !== !!params.ref) {
+      throw new Error("Both element and ref must be provided or neither.");
+    }
+
     const page = await context.getActivePage();
     if (!page) {
       throw new Error("No active page found for screenshot");
@@ -407,15 +407,12 @@ const screenshot = defineTool<typeof screenshotSchema>({
       `screenshot-${Date.now()}.${fileType}`
     );
 
-    const baseOptions: Omit<
-      Parameters<Page["screenshot"]>[0],
-      "type" | "quality" | "path"
-    > = {
+    const baseOptions: PageScreenshotOptions = {
       scale: "css",
       timeout: 15000, // Kept existing timeout
     };
 
-    let options: Parameters<Page["screenshot"]>[0];
+    let options: PageScreenshotOptions;
 
     if (fileType === "jpeg") {
       options = {
@@ -499,4 +496,4 @@ export async function generateLocator(locator: Locator): Promise<string> {
   return (locator as any)._generateLocatorString();
 }
 
-export default [snapshot, click, drag, hover, type, selectOption, screenshot];
+export default [snapshot, click, drag, hover, type, selectOption, screenshot];