fix: Address issues with session management and code generation of new session tool

Winify · Winify · commit 98642537a76b · 2026-03-30T09:33:42.000+02:00
- Simplify session metadata handling by directly using merged capabilities.
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
@@ -8,6 +8,9 @@ on:
     branches:
       - main
 
+permissions:
+  contents: read
+
 jobs:
   lint:
     runs-on: ubuntu-latest
diff --git a/.gitignore b/.gitignore
@@ -4,4 +4,5 @@
 node_modules
 lib
 
+settings.local.json
 *.tgz
diff --git a/src/recording/code-generator.ts b/src/recording/code-generator.ts
@@ -31,11 +31,13 @@ function generateStep(step: RecordedStep, history: SessionHistory): string {
 
   const p = step.params;
   switch (step.tool) {
-    case 'start_browser': {
-      const nav = p.navigationUrl ? `\nawait browser.url('${escapeStr(p.navigationUrl)}');` : '';
-      return `const browser = await remote({\n  capabilities: ${indentJson(history.capabilities)}\n});${nav}`;
-    }
-    case 'start_app_session': {
+    case 'start_session': {
+      const platform = p.platform as string;
+      if (platform === 'browser') {
+        const nav = p.navigationUrl ? `\nawait browser.url('${escapeStr(p.navigationUrl)}');` : '';
+        return `const browser = await remote({\n  capabilities: ${indentJson(history.capabilities)}\n});${nav}`;
+      }
+      // Mobile (ios/android)
       const config: Record<string, unknown> = {
         protocol: 'http',
         hostname: history.appiumConfig?.hostname ?? 'localhost',
@@ -45,10 +47,8 @@ function generateStep(step: RecordedStep, history: SessionHistory): string {
       };
       return `const browser = await remote(${indentJson(config)});`;
     }
-    case 'attach_browser': {
-      const nav = p.navigationUrl ? `\nawait browser.url('${escapeStr(p.navigationUrl)}');` : '';
-      return `const browser = await remote({\n  capabilities: ${indentJson(history.capabilities)}\n});${nav}`;
-    }
+    case 'close_session':
+      return 'await browser.deleteSession();';
     case 'navigate':
       return `await browser.url('${escapeStr(p.url)}');`;
     case 'click_element':
@@ -78,5 +78,5 @@ function generateStep(step: RecordedStep, history: SessionHistory): string {
 
 export function generateCode(history: SessionHistory): string {
   const steps = history.steps.map(step => generateStep(step, history)).join('\n');
-  return `import { remote } from 'webdriverio';\n\n${steps}\n\nawait browser.deleteSession();`;
+  return `import { remote } from 'webdriverio';\n\n${steps}`;
 }
diff --git a/src/session/lifecycle.ts b/src/session/lifecycle.ts
@@ -27,16 +27,30 @@ export function registerSession(
   historyEntry: SessionHistory,
 ): void {
   const state = getState();
-  if (state.currentSession && state.currentSession !== sessionId) {
+  const oldSessionId = state.currentSession;
+  if (oldSessionId && oldSessionId !== sessionId) {
     handleSessionTransition(sessionId);
   }
   state.browsers.set(sessionId, browser);
   state.sessionMetadata.set(sessionId, metadata);
   state.sessionHistory.set(sessionId, historyEntry);
   state.currentSession = sessionId;
+
+  // If there was a previous session, terminate it to prevent orphaning
+  if (oldSessionId && oldSessionId !== sessionId) {
+    const oldBrowser = state.browsers.get(oldSessionId);
+    if (oldBrowser) {
+      // Fire and forget — don't block registration on close
+      oldBrowser.deleteSession().catch(() => {
+        // Ignore errors during force-close of orphaned session
+      });
+      state.browsers.delete(oldSessionId);
+      state.sessionMetadata.delete(oldSessionId);
+    }
+  }
 }
 
-export async function closeSession(sessionId: string, detach: boolean, isAttached: boolean): Promise<void> {
+export async function closeSession(sessionId: string, detach: boolean, isAttached: boolean, force?: boolean): Promise<void> {
   const state = getState();
   const browser = state.browsers.get(sessionId);
   if (!browser) return;
@@ -46,12 +60,18 @@ export async function closeSession(sessionId: string, detach: boolean, isAttache
     history.endedAt = new Date().toISOString();
   }
 
-  // Only terminate the WebDriver session if we created it (not attached/borrowed)
-  if (!detach && !isAttached) {
+  // Terminate the WebDriver session if:
+  // - force is true (override), OR
+  // - detach is false AND isAttached is false (normal close)
+  if (force || (!detach && !isAttached)) {
     await browser.deleteSession();
   }
 
   state.browsers.delete(sessionId);
   state.sessionMetadata.delete(sessionId);
-  state.currentSession = null;
+
+  // Only clear currentSession if it matches the session being closed
+  if (state.currentSession === sessionId) {
+    state.currentSession = null;
+  }
 }
diff --git a/src/tools/session.tool.ts b/src/tools/session.tool.ts
@@ -166,15 +166,17 @@ async function startBrowserSession(args: StartSessionArgs): Promise<CallToolResu
   const wdioBrowser = await remote({ capabilities: mergedCapabilities });
   const { sessionId } = wdioBrowser;
 
-  registerSession(sessionId, wdioBrowser, {
+  const sessionMetadata: SessionMetadata = {
     type: 'browser',
-    capabilities: wdioBrowser.capabilities as Record<string, unknown>,
+    capabilities: mergedCapabilities,
     isAttached: false,
-  }, {
+  };
+
+  registerSession(sessionId, wdioBrowser, sessionMetadata, {
     sessionId,
     type: 'browser',
     startedAt: new Date().toISOString(),
-    capabilities: wdioBrowser.capabilities as Record<string, unknown>,
+    capabilities: mergedCapabilities,
     steps: [],
   });
 
@@ -267,41 +269,36 @@ async function attachBrowserSession(args: StartSessionArgs): Promise<CallToolRes
   await waitForCDP(host, port);
   const { activeTabUrl, allTabUrls } = await closeStaleMappers(host, port);
 
+  const capabilities = {
+    browserName: 'chrome',
+    unhandledPromptBehavior: 'dismiss',
+    webSocketUrl: false,
+    'goog:chromeOptions': {
+      debuggerAddress: `${host}:${port}`,
+    },
+  };
+
   const browser = await remote({
     connectionRetryTimeout: 30000,
     connectionRetryCount: 3,
-    capabilities: {
-      browserName: 'chrome',
-      unhandledPromptBehavior: 'dismiss',
-      webSocketUrl: false,
-      'goog:chromeOptions': {
-        debuggerAddress: `${host}:${port}`,
-      },
-    },
+    capabilities,
   });
 
   const { sessionId } = browser;
-  registerSession(
+
+  const sessionMetadata: SessionMetadata = {
+    type: 'browser',
+    capabilities,
+    isAttached: true,
+  };
+
+  registerSession(sessionId, browser, sessionMetadata, {
     sessionId,
-    browser,
-    {
-      type: 'browser',
-      capabilities: browser.capabilities as Record<string, unknown>,
-      isAttached: true,
-    },
-    {
-      sessionId,
-      type: 'browser',
-      startedAt: new Date().toISOString(),
-      capabilities: {
-        browserName: 'chrome',
-        'goog:chromeOptions': {
-          debuggerAddress: `${host}:${port}`,
-        },
-      },
-      steps: [],
-    },
-  );
+    type: 'browser',
+    startedAt: new Date().toISOString(),
+    capabilities,
+    steps: [],
+  });
 
   if (navigationUrl) {
     await browser.url(navigationUrl);
@@ -341,11 +338,17 @@ export const closeSessionTool: ToolCallback = async (args: { detach?: boolean }
     const sessionId = state.currentSession;
     const metadata = state.sessionMetadata.get(sessionId);
 
-    const effectiveDetach = args.detach || !!metadata?.isAttached;
-    await closeSession(sessionId, args.detach ?? false, !!metadata?.isAttached);
+    const isAttached = !!metadata?.isAttached;
+    const detach = args.detach ?? false;
+
+    // Determine if this is a force-close (auto-detached session + explicit close)
+    const force = !detach && isAttached;
+
+    const effectiveDetach = detach || isAttached;
+    await closeSession(sessionId, detach, isAttached, force);
 
-    const action = effectiveDetach ? 'detached from' : 'closed';
-    const note = args.detach && !metadata?.isAttached
+    const action = effectiveDetach && !force ? 'detached from' : 'closed';
+    const note = detach && !isAttached
       ? '\nNote: Session will remain active on Appium server.'
       : '';
 
diff --git a/tests/recording/code-generator.test.ts b/tests/recording/code-generator.test.ts
@@ -5,8 +5,8 @@ import type { SessionHistory, RecordedStep } from '../../src/types/recording';
 
 const START_BROWSER_STEP: RecordedStep = {
   index: 1,
-  tool: 'start_browser',
-  params: { browser: 'chrome', headless: true, windowWidth: 1920, windowHeight: 1080 },
+  tool: 'start_session',
+  params: { platform: 'browser', browser: 'chrome', headless: true, windowWidth: 1920, windowHeight: 1080 },
   status: 'ok',
   durationMs: 0,
   timestamp: '2026-01-01T00:00:00.000Z',
@@ -35,14 +35,13 @@ function makeHistory(steps: Partial<RecordedStep>[]): SessionHistory {
 }
 
 describe('generateCode - header', () => {
-  it('wraps output in remote() setup and deleteSession', () => {
+  it('wraps output in remote() setup', () => {
     const code = generateCode(makeHistory([]));
     expect(code).toContain("import { remote } from 'webdriverio';");
-    expect(code).toContain('await browser.deleteSession();');
     expect(code).toContain('browserName');
   });
 
-  it('generates start_browser using history.capabilities, not reconstructed from params', () => {
+  it('generates start_session (browser) using history.capabilities', () => {
     const history: SessionHistory = {
       sessionId: 'caps-123',
       type: 'browser',
@@ -54,8 +53,8 @@ describe('generateCode - header', () => {
       },
       steps: [{
         index: 1,
-        tool: 'start_browser',
-        params: { browser: 'chrome', headless: true },
+        tool: 'start_session',
+        params: { platform: 'browser', browser: 'chrome', headless: true },
         status: 'ok',
         durationMs: 100,
         timestamp: '2026-01-01T00:00:00.000Z',
@@ -64,10 +63,10 @@ describe('generateCode - header', () => {
     const code = generateCode(history);
     expect(code).toContain('const browser = await remote(');
     expect(code).toContain('"browserName": "chrome"');
-    expect(code).toContain('--custom-flag'); // only present in history.capabilities
+    expect(code).toContain('--custom-flag');
   });
 
-  it('generates attach_browser using history.capabilities', () => {
+  it('generates start_session (attach mode) using history.capabilities', () => {
     const history: SessionHistory = {
       sessionId: 'attach-123',
       type: 'browser',
@@ -78,8 +77,8 @@ describe('generateCode - header', () => {
       },
       steps: [{
         index: 1,
-        tool: 'attach_browser',
-        params: { port: 9222, host: 'localhost', userDataDir: '/tmp/chrome-debug' },
+        tool: 'start_session',
+        params: { platform: 'browser', attach: true, port: 9222, host: 'localhost', userDataDir: '/tmp/chrome-debug' },
         status: 'ok',
         durationMs: 100,
         timestamp: '2026-01-01T00:00:00.000Z',
@@ -91,16 +90,16 @@ describe('generateCode - header', () => {
     expect(code).toContain('--user-data-dir=/tmp/chrome-debug');
   });
 
-  it('appends browser.url() when navigationUrl is set on start_browser', () => {
+  it('appends browser.url() when navigationUrl is set on start_session (browser)', () => {
     const history: SessionHistory = {
       sessionId: 'nav-123',
       type: 'browser',
       startedAt: '2026-01-01T00:00:00.000Z',
       capabilities: { browserName: 'chrome' },
       steps: [{
         index: 1,
-        tool: 'start_browser',
-        params: { browser: 'chrome', headless: false, windowWidth: 1920, windowHeight: 1080, navigationUrl: 'https://github.com/login' },
+        tool: 'start_session',
+        params: { platform: 'browser', browser: 'chrome', headless: false, windowWidth: 1920, windowHeight: 1080, navigationUrl: 'https://github.com/login' },
         status: 'ok',
         durationMs: 0,
         timestamp: '2026-01-01T00:00:00.000Z',
@@ -110,7 +109,7 @@ describe('generateCode - header', () => {
     expect(code).toContain("await browser.url('https://github.com/login');");
   });
 
-  it('generates start_app_session using history.appiumConfig for connection config', () => {
+  it('generates start_session (mobile) using history.appiumConfig for connection config', () => {
     const history: SessionHistory = {
       sessionId: 'app-123',
       type: 'android',
@@ -123,15 +122,15 @@ describe('generateCode - header', () => {
       appiumConfig: { hostname: '127.0.0.1', port: 4723, path: '/' },
       steps: [{
         index: 1,
-        tool: 'start_app_session',
-        params: { platform: 'Android', deviceName: 'emulator-5554' }, // no appiumHost in params
+        tool: 'start_session',
+        params: { platform: 'android', deviceName: 'emulator-5554' },
         status: 'ok',
         durationMs: 100,
         timestamp: '2026-01-01T00:00:00.000Z',
       }],
     };
     const code = generateCode(history);
-    expect(code).toContain('"hostname": "127.0.0.1"'); // from history.appiumConfig, not params fallback
+    expect(code).toContain('"hostname": "127.0.0.1"');
     expect(code).toContain('"port": 4723');
     expect(code).toContain('"platformName": "Android"');
   });
diff --git a/tests/session/lifecycle.test.ts b/tests/session/lifecycle.test.ts
@@ -5,7 +5,7 @@ import { closeSession, registerSession } from '../../src/session/lifecycle';
 import type { SessionHistory } from '../../src/types/recording';
 
 function makeBrowser(overrides: Record<string, unknown> = {}) {
-  return { deleteSession: vi.fn(), ...overrides } as unknown as WebdriverIO.Browser;
+  return { deleteSession: vi.fn().mockResolvedValue(undefined), ...overrides } as unknown as WebdriverIO.Browser;
 }
 
 beforeEach(() => {
diff --git a/tests/tools/close-session.test.ts b/tests/tools/close-session.test.ts
@@ -41,15 +41,27 @@ describe('close_session', () => {
     expect(mockDeleteSession).toHaveBeenCalledOnce();
   });
 
-  it('skips deleteSession when isAttached is true', async () => {
+  it('calls deleteSession when isAttached is true (force-close)', async () => {
     setupSession('sess-attached', true);
     await callClose({});
-    expect(mockDeleteSession).not.toHaveBeenCalled();
+    expect(mockDeleteSession).toHaveBeenCalledOnce();
   });
 
-  it('returns "detached from" message when isAttached is true and detach is false', async () => {
+  it('returns "closed" message when isAttached is true (force-close)', async () => {
     setupSession('sess-attached', true);
     const result = await callClose({});
+    expect(result.content[0].text).toContain('closed');
+  });
+
+  it('skips deleteSession when detach is explicitly true', async () => {
+    setupSession('sess-attached', true);
+    await callClose({ detach: true });
+    expect(mockDeleteSession).not.toHaveBeenCalled();
+  });
+
+  it('returns "detached from" message when detach is explicitly true', async () => {
+    setupSession('sess-attached', true);
+    const result = await callClose({ detach: true });
     expect(result.content[0].text).toContain('detached from');
   });
 

Original file line number	Diff line number	Diff line change
`@@ -5,7 +5,7 @@ import { closeSession, registerSession } from '../../src/session/lifecycle';`
`5`	`5`	`import type { SessionHistory } from '../../src/types/recording';`
`6`	`6`
`7`	`7`	`function makeBrowser(overrides: Record<string, unknown> = {}) {`
`8`		`- return { deleteSession: vi.fn(), ...overrides } as unknown as WebdriverIO.Browser;`
	`8`	`+ return { deleteSession: vi.fn().mockResolvedValue(undefined), ...overrides } as unknown as WebdriverIO.Browser;`
`9`	`9`	`}`
`10`	`10`
`11`	`11`	`beforeEach(() => {`