@amaster.ai/pi-computer-use 0.1.2-beta.8 → 0.1.2-beta.80

package/dist/index.js

@@ -1,12 +1,68 @@
+import { lstat, realpath } from 'node:fs/promises';
+import path from 'node:path';
+import { isProjectTrusted } from '@amaster.ai/pi-shared/settings';
+import { StringEnum } from '@earendil-works/pi-ai';
 import { Type } from 'typebox';
 import { loadConfigFromFile, resolveConfig } from './config.js';
-import { CuaDriverClient } from './mcp-client.js';
+import toolManifest from './generated/cua-driver-tools.js';
+import { CuaDriverClient, waitForPromise } from './mcp-client.js';
+import { CORE_TOOLS, TOOL_GROUP_NAMES, TOOL_GROUPS } from './tool-groups.js';
+import { toPiToolResult } from './tool-result.js';
 import { createPiVisionCaller } from './vision.js';
+export { CuaDriverClient, resolveBundledTarget, resolveDriverLayout, resolveUnixSocketPath, } from './mcp-client.js';
 export { loadConfigFromFile, resolveConfig };
 const TOOL_PREFIX = 'computer_use_';
-const PLATFORM = process.platform;
+const HIGH_RISK_TOOLS = new Set([
+    'browser_download',
+    'browser_prepare',
+    'browser_set_input_files',
+    'clipboard_read',
+    'clipboard_write',
+    'install_ffmpeg',
+    'kill_app',
+    'replay_trajectory',
+    'start_recording',
+]);
+function isInside(relativePath) {
+    return (relativePath === '' ||
+        (relativePath !== '..' &&
+            !relativePath.startsWith(`..${path.sep}`) &&
+            !path.isAbsolute(relativePath)));
+}
+async function approvedRecordingDir(value, cwd) {
+    const root = path.resolve(cwd);
+    const outputDir = path.resolve(cwd, String(value ?? ''));
+    const lexicalRelative = path.relative(root, outputDir);
+    if (!isInside(lexicalRelative))
+        return undefined;
+    try {
+        const canonicalRoot = await realpath(root);
+        let existing = outputDir;
+        const suffix = [];
+        for (;;) {
+            const info = await lstat(existing).catch(() => undefined);
+            if (info) {
+                if (!info.isDirectory() && !info.isSymbolicLink())
+                    return undefined;
+                if (existing === outputDir && info.isSymbolicLink())
+                    return undefined;
+                const canonical = path.resolve(await realpath(existing), ...suffix);
+                const canonicalRelative = path.relative(canonicalRoot, canonical);
+                return isInside(canonicalRelative) ? canonical : undefined;
+            }
+            const parent = path.dirname(existing);
+            if (parent === existing)
+                return undefined;
+            suffix.unshift(path.basename(existing));
+            existing = parent;
+        }
+    }
+    catch {
+        return undefined;
+    }
+}
 function permissionHint() {
-    switch (PLATFORM) {
+    switch (process.platform) {
         case 'darwin':
             return 'Check that Accessibility and Screen Recording permissions are granted in System Settings → Privacy & Security.';
         case 'win32':
@@ -16,7 +72,7 @@ function permissionHint() {
     }
 }
 function accessibilityHint() {
-    switch (PLATFORM) {
+    switch (process.platform) {
         case 'darwin':
             return 'Accessibility permission not granted. The user needs to enable it in System Settings → Privacy & Security → Accessibility, then restart the app.';
         case 'win32':
@@ -26,7 +82,7 @@ function accessibilityHint() {
     }
 }
 function screenRecordingHint() {
-    switch (PLATFORM) {
+    switch (process.platform) {
         case 'darwin':
             return 'Screen Recording permission not granted. The user needs to enable it in System Settings → Privacy & Security → Screen & System Audio Recording, then restart the app.';
         case 'win32':
@@ -35,380 +91,257 @@ function screenRecordingHint() {
             return 'Screen capture failed. Check that the compositor allows screen capture (PipeWire portal or X11 access).';
     }
 }
-const EXCLUDED_TOOLS = new Set([
-    'set_agent_cursor_enabled',
-    'set_agent_cursor_motion',
-    'set_agent_cursor_style',
-    'get_agent_cursor_state',
-    'set_recording',
-    'get_recording_state',
-    'replay_trajectory',
-    'check_permissions',
-    'get_config',
-    'set_config',
-    'move_cursor',
-    'zoom',
-    'type_text_chars',
-    'page',
-    'browser_eval',
-    'screenshot',
-]);
-const FALLBACK_TOOLS = [
-    {
-        name: 'click',
-        description: 'Left-click against a target pid via element_index or x/y coordinates',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                pid: { type: 'integer', description: 'Target process ID' },
-                x: { type: 'number', description: 'Window-local screenshot X coordinate' },
-                y: { type: 'number', description: 'Window-local screenshot Y coordinate' },
-                element_index: { type: 'integer', description: 'Element index from last get_window_state' },
-                window_id: { type: 'integer', description: 'Target window ID. Required for element_index' },
-                action: {
-                    type: 'string',
-                    description: 'AX action: press, show_menu, pick, confirm, cancel, open',
-                },
-                modifier: {
-                    type: 'array',
-                    items: { type: 'string' },
-                    description: 'Modifier keys: cmd, shift, option/alt, ctrl',
-                },
-                from_zoom: {
-                    type: 'boolean',
-                    description: 'When true, x/y are in last zoom image coordinates',
-                },
-            },
-            required: ['pid'],
-        },
-    },
-    {
-        name: 'double_click',
-        description: 'Double-click at x/y or on an AX element via element_index',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                pid: { type: 'integer' },
-                x: { type: 'number', description: 'Screen X coordinate (pixel path)' },
-                y: { type: 'number', description: 'Screen Y coordinate (pixel path)' },
-                element_index: { type: 'integer', description: 'Element index from last get_window_state' },
-                window_id: {
-                    type: 'integer',
-                    description: 'CGWindowID. Required when element_index is used',
-                },
-            },
-            required: ['pid'],
-        },
-    },
-    {
-        name: 'right_click',
-        description: 'Right-click against a target pid via element_index or x/y coordinates',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                pid: { type: 'integer', description: 'Target process ID' },
-                x: { type: 'number', description: 'X in window-local screenshot pixels' },
-                y: { type: 'number', description: 'Y in window-local screenshot pixels' },
-                element_index: {
-                    type: 'integer',
-                    description: 'Element index from last get_window_state. Routes through AXShowMenu',
-                },
-                window_id: {
-                    type: 'integer',
-                    description: 'CGWindowID. Required when element_index is used',
-                },
-                modifier: {
-                    type: 'array',
-                    items: { type: 'string' },
-                    description: 'Modifier keys held during the right-click (pixel path only)',
-                },
-            },
-            required: ['pid'],
-        },
-    },
-    {
-        name: 'type_text',
-        description: 'Insert text into the target pid via AX or CGEvent fallback',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                pid: { type: 'integer', description: 'Target process ID' },
-                text: { type: 'string', description: 'Text to insert at the target cursor' },
-                element_index: { type: 'integer', description: 'Element index from last get_window_state' },
-                window_id: {
-                    type: 'integer',
-                    description: 'CGWindowID. Required when element_index is used',
-                },
-                delay_ms: {
-                    type: 'integer',
-                    minimum: 0,
-                    maximum: 200,
-                    description: 'Milliseconds between characters in CGEvent fallback. Default 30',
-                },
-            },
-            required: ['pid', 'text'],
-        },
-    },
-    {
-        name: 'press_key',
-        description: 'Press and release a single key, delivered to the target pid',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                pid: { type: 'integer' },
-                key: {
-                    type: 'string',
-                    description: 'Key name: return, tab, escape, up, down, left, right, space, delete, etc.',
-                },
-                modifiers: {
-                    type: 'array',
-                    items: { type: 'string' },
-                    description: 'Modifier keys: cmd, shift, option/alt, ctrl, fn',
-                },
-                element_index: { type: 'integer' },
-                window_id: { type: 'integer' },
-            },
-            required: ['pid', 'key'],
-        },
-    },
-    {
-        name: 'hotkey',
-        description: 'Press a combination of keys simultaneously, e.g. ["cmd", "c"] for Copy',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                pid: { type: 'integer', description: 'Target process ID' },
-                keys: {
-                    type: 'array',
-                    items: { type: 'string' },
-                    minItems: 2,
-                    description: 'Modifier(s) and one non-modifier key, e.g. ["cmd", "c"]',
-                },
-                window_id: {
-                    type: 'integer',
-                    description: 'When set, uses NSMenu path for native menu key dispatch',
-                },
-            },
-            required: ['pid', 'keys'],
-        },
-    },
-    {
-        name: 'scroll',
-        description: 'Scroll the target pid focused region by synthesized keystrokes',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                pid: { type: 'integer' },
-                direction: {
-                    type: 'string',
-                    enum: ['up', 'down', 'left', 'right'],
-                    description: 'Scroll direction',
-                },
-                amount: {
-                    type: 'integer',
-                    minimum: 1,
-                    maximum: 50,
-                    description: 'Number of keystroke repetitions. Default: 3',
-                },
-                by: {
-                    type: 'string',
-                    enum: ['line', 'page'],
-                    description: 'Scroll granularity. Default: line',
-                },
-                element_index: { type: 'integer' },
-                window_id: { type: 'integer' },
-            },
-            required: ['pid', 'direction'],
-        },
-    },
-    {
-        name: 'drag',
-        description: 'Press-drag-release gesture from one point to another in window-local pixels',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                pid: { type: 'integer', description: 'Target process ID' },
-                from_x: { type: 'number', description: 'Drag-start X in window-local screenshot pixels' },
-                from_y: { type: 'number', description: 'Drag-start Y in window-local screenshot pixels' },
-                to_x: { type: 'number', description: 'Drag-end X in window-local screenshot pixels' },
-                to_y: { type: 'number', description: 'Drag-end Y in window-local screenshot pixels' },
-                button: {
-                    type: 'string',
-                    enum: ['left', 'right', 'middle'],
-                    description: 'Mouse button. Default: left',
-                },
-                duration_ms: {
-                    type: 'integer',
-                    minimum: 0,
-                    maximum: 10000,
-                    description: 'Duration of drag path. Default: 500',
-                },
-                steps: {
-                    type: 'integer',
-                    minimum: 1,
-                    maximum: 200,
-                    description: 'Intermediate drag events. Default: 20',
-                },
-                modifier: {
-                    type: 'array',
-                    items: { type: 'string' },
-                    description: 'Modifier keys held across the gesture',
-                },
-                window_id: { type: 'integer' },
-                from_zoom: { type: 'boolean' },
-            },
-            required: ['pid', 'from_x', 'from_y', 'to_x', 'to_y'],
-        },
-    },
-    {
-        name: 'set_value',
-        description: 'Set a value on a UI element (popups, sliders, steppers, date pickers)',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                pid: { type: 'integer' },
-                window_id: {
-                    type: 'integer',
-                    description: 'CGWindowID for the window whose get_window_state produced the element_index',
-                },
-                element_index: { type: 'integer' },
-                value: {
-                    type: 'string',
-                    description: 'New value. AX will coerce to the element native type',
-                },
-            },
-            required: ['pid', 'window_id', 'element_index', 'value'],
-        },
-    },
-    {
-        name: 'get_screen_size',
-        description: 'Return the logical size of the main display in points plus backing scale factor',
-        inputSchema: { type: 'object', additionalProperties: false, properties: {} },
-    },
-    {
-        name: 'get_cursor_position',
-        description: 'Return the current mouse cursor position in screen points',
-        inputSchema: { type: 'object', additionalProperties: false, properties: {} },
-    },
-    {
-        name: 'get_accessibility_tree',
-        description: 'Return a lightweight desktop snapshot: running apps and visible windows with bounds and z-order',
-        inputSchema: { type: 'object', additionalProperties: false, properties: {} },
-    },
-    {
-        name: 'get_window_state',
-        description: 'Walk an app AX tree and return a Markdown rendering of its UI with actionable element indices',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                pid: { type: 'integer', description: 'Target process ID' },
-                window_id: { type: 'integer', description: 'Target window ID from list_windows' },
-                query: { type: 'string', description: 'Case-insensitive filter for tree_markdown' },
-                capture_mode: {
-                    type: 'string',
-                    enum: ['som', 'vision', 'ax'],
-                    description: 'som=AX+screenshot (default), vision=screenshot only, ax=AX only',
-                },
-            },
-            required: ['pid', 'window_id'],
-        },
-    },
-    {
-        name: 'list_windows',
-        description: 'List all layer-0 top-level windows known to WindowServer',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                pid: { type: 'integer', description: 'Optional pid filter' },
-                on_screen_only: {
-                    type: 'boolean',
-                    description: 'When true, drop windows not on current Space. Default false',
-                },
-            },
-        },
-    },
-    {
-        name: 'list_apps',
-        description: 'List macOS apps (running and installed) with state flags, pid, bundle_id',
-        inputSchema: { type: 'object', additionalProperties: false, properties: {} },
-    },
-    {
-        name: 'launch_app',
-        description: 'Launch a macOS app in the background without stealing focus',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: {
-                bundle_id: {
-                    type: 'string',
-                    description: 'App bundle identifier, e.g. com.apple.calculator. Preferred over name',
-                },
-                name: {
-                    type: 'string',
-                    description: 'App display name. Used only when bundle_id is absent',
-                },
-                urls: {
-                    type: 'array',
-                    items: { type: 'string' },
-                    description: 'File paths or URLs to open with the app',
-                },
-                creates_new_application_instance: {
-                    type: 'boolean',
-                    description: 'Force a new app instance even if already running',
-                },
-            },
-        },
-    },
-    {
-        name: 'kill_app',
-        description: 'Force-terminate a process by pid (kill -9 equivalent)',
-        inputSchema: {
-            type: 'object',
-            additionalProperties: false,
-            properties: { pid: { type: 'integer', description: 'PID of the process to terminate' } },
-            required: ['pid'],
-        },
-    },
-];
 export default function computerUseExtension(pi) {
     let config;
     let client;
-    let connected = false;
-    async function ensureConnected() {
-        if (!client)
+    let sessionAbortController;
+    let macPermissionPromise;
+    const approvedLaunchApprovalKeys = new Set();
+    const driverToolNames = new Set();
+    const TOOLS_META_TOOL = `${TOOL_PREFIX}tools`;
+    async function ensureConnected(signal, requestMacPermissions = true) {
+        signal?.throwIfAborted();
+        if (!client) {
+            if (!config)
+                throw new Error('pi-computer-use: session not started');
+            client = new CuaDriverClient(config);
+        }
+        const sessionClient = client;
+        const sessionSignal = sessionAbortController?.signal;
+        if (!sessionSignal)
             throw new Error('pi-computer-use: session not started');
-        if (!connected) {
-            await client.connect();
-            connected = true;
+        sessionSignal.throwIfAborted();
+        await sessionClient.ensureReady(signal);
+        if (process.platform !== 'darwin' || !requestMacPermissions)
+            return;
+        sessionSignal.throwIfAborted();
+        if (!macPermissionPromise) {
+            const permissionPromise = sessionClient
+                .callTool('check_permissions', { prompt: true }, sessionSignal)
+                .then(() => undefined, () => {
+                if (!sessionSignal.aborted) {
+                    console.error('[pi-computer-use] macOS permission probe failed; requested tool will continue');
+                }
+            });
+            macPermissionPromise = permissionPromise;
+        }
+        await waitForPromise(macPermissionPromise, signal);
+        sessionSignal.throwIfAborted();
+    }
+    async function confirmToolCall(toolName, params, ctx) {
+        const requestedRecordingDir = params.output_dir;
+        if (toolName === 'launch_app' && config?.confirmAppLaunch !== false) {
+            const target = String(params.bundle_id ??
+                params.aumid ??
+                params.name ??
+                params.app_name ??
+                params.path ??
+                params.command ??
+                'unknown app');
+            const approvalKey = JSON.stringify(Object.fromEntries(Object.entries(params).sort(([a], [b]) => a.localeCompare(b))));
+            if (approvedLaunchApprovalKeys.has(approvalKey))
+                return true;
+            if (!ctx.hasUI)
+                return false;
+            const approved = await ctx.ui.confirm('Allow computer use?', `Allow pi-computer-use to launch and control ${target} for this session?`);
+            if (approved)
+                approvedLaunchApprovalKeys.add(approvalKey);
+            return approved;
+        }
+        if (toolName === 'start_recording') {
+            const outputDir = await approvedRecordingDir(requestedRecordingDir, ctx.cwd);
+            if (!outputDir)
+                return false;
+            params.output_dir = outputDir;
+        }
+        if (toolName === 'start_recording' ||
+            (HIGH_RISK_TOOLS.has(toolName) && config?.confirmDangerousActions !== false)) {
+            if (!ctx.hasUI)
+                return false;
+            const approved = await ctx.ui.confirm('Confirm high-risk computer action', `Allow ${toolName} with arguments ${JSON.stringify(params)}?`);
+            if (!approved)
+                return false;
+            if (toolName === 'start_recording') {
+                const outputDir = await approvedRecordingDir(requestedRecordingDir, ctx.cwd);
+                if (!outputDir || outputDir !== params.output_dir)
+                    return false;
+                params.output_dir = outputDir;
+            }
+            return true;
+        }
+        return true;
+    }
+    function formatGroupList() {
+        const active = new Set(pi.getActiveTools());
+        const lines = ['Computer-use tool groups:'];
+        for (const [name, group] of Object.entries(TOOL_GROUPS)) {
+            const registered = group.tools.filter((tool) => driverToolNames.has(`${TOOL_PREFIX}${tool}`));
+            const activeCount = registered.filter((tool) => active.has(`${TOOL_PREFIX}${tool}`)).length;
+            const suffix = registered.length === 0
+                ? 'unavailable until the driver is registered'
+                : activeCount === 0
+                    ? 'inactive'
+                    : activeCount === registered.length
+                        ? 'active'
+                        : `${activeCount}/${registered.length} active`;
+            lines.push(`- ${name} (${group.tools.length} tools, ${suffix}): ${group.summary}`);
         }
+        lines.push(`Activate a group with ${TOOLS_META_TOOL} group=<name> or /computer-use-tools <name>.`);
+        return lines.join('\n');
+    }
+    function activateGroup(group) {
+        const definition = TOOL_GROUPS[group];
+        const active = pi.getActiveTools();
+        const wanted = (definition?.tools ?? [])
+            .map((tool) => `${TOOL_PREFIX}${tool}`)
+            .filter((name) => driverToolNames.has(name));
+        const toAdd = wanted.filter((name) => !active.includes(name));
+        if (toAdd.length > 0)
+            pi.setActiveTools([...active, ...toAdd]);
+        return {
+            activated: wanted,
+            total: definition?.tools.length ?? 0,
+            alreadyActive: toAdd.length === 0 && wanted.length > 0,
+        };
+    }
+    function applyToolProfile() {
+        if (config?.toolProfile !== 'core')
+            return;
+        const core = new Set(CORE_TOOLS);
+        pi.setActiveTools(pi
+            .getActiveTools()
+            .filter((name) => !driverToolNames.has(name) || core.has(name.slice(TOOL_PREFIX.length))));
+    }
+    function registerToolGroupSurface() {
+        const groupSummary = Object.entries(TOOL_GROUPS)
+            .map(([name, group]) => `${name} (${group.summary})`)
+            .join(', ');
+        pi.registerTool({
+            name: TOOLS_META_TOOL,
+            label: TOOLS_META_TOOL,
+            description: `List and activate additional computer-use tool groups: ${groupSummary}. Call without arguments to list groups and their activation status.`,
+            parameters: Type.Object({
+                group: Type.Optional(StringEnum(TOOL_GROUP_NAMES, {
+                    description: 'Tool group to activate. Omit to list available groups.',
+                })),
+            }),
+            promptSnippet: 'List or activate extra computer-use tool groups when the core tools are insufficient',
+            promptGuidelines: [
+                `Only the core computer-use tools are active by default. When a task needs browser automation, recording, escalated sessions, cursor overlay, window management, clipboard, or diagnostics capabilities, call ${TOOLS_META_TOOL} with the matching group first — the group's tools become available on the next turn.`,
+            ],
+            async execute(_toolCallId, params) {
+                const group = typeof params.group === 'string' ? params.group : undefined;
+                if (!group) {
+                    return {
+                        content: [{ type: 'text', text: formatGroupList() }],
+                        details: undefined,
+                    };
+                }
+                if (driverToolNames.size === 0) {
+                    return {
+                        content: [
+                            {
+                                type: 'text',
+                                text: `Driver tools are not registered yet — connect the driver first with ${TOOL_PREFIX}connect, then activate the "${group}" group again.`,
+                            },
+                        ],
+                        details: undefined,
+                        isError: true,
+                    };
+                }
+                const { activated, total, alreadyActive } = activateGroup(group);
+                if (activated.length === 0) {
+                    return {
+                        content: [
+                            {
+                                type: 'text',
+                                text: `Group "${group}" has no tools registered on this platform's driver.`,
+                            },
+                        ],
+                        details: undefined,
+                        isError: true,
+                    };
+                }
+                return {
+                    content: [
+                        {
+                            type: 'text',
+                            text: `${alreadyActive ? 'Already active' : 'Activated'} group "${group}" (${activated.length}/${total} tools): ${activated.join(', ')}. These tools are available from the next turn.`,
+                        },
+                    ],
+                    details: undefined,
+                };
+            },
+        });
+        pi.registerCommand('computer-use-tools', {
+            description: 'List or activate computer-use tool groups',
+            getArgumentCompletions: (argumentPrefix) => {
+                const prefix = argumentPrefix.trim().toLowerCase();
+                return Object.entries(TOOL_GROUPS)
+                    .filter(([name]) => name.startsWith(prefix))
+                    .map(([name, group]) => ({ value: name, label: name, description: group.summary }));
+            },
+            async handler(args, ctx) {
+                const group = String(args ?? '')
+                    .trim()
+                    .toLowerCase();
+                if (!group) {
+                    ctx.ui.notify(formatGroupList(), 'info');
+                    return;
+                }
+                if (!TOOL_GROUP_NAMES.includes(group)) {
+                    ctx.ui.notify(`pi-computer-use: unknown group "${group}". Available: ${TOOL_GROUP_NAMES.join(', ')}`, 'warning');
+                    return;
+                }
+                const { activated, alreadyActive } = activateGroup(group);
+                ctx.ui.notify(activated.length === 0
+                    ? `pi-computer-use: group "${group}" has no registered tools (driver not connected?).`
+                    : `pi-computer-use: ${alreadyActive ? 'already active' : 'activated'} "${group}" (${activated.length} tools).`, 'info');
+            },
+        });
     }
     function registerTools(tools) {
         for (const tool of tools) {
-            if (EXCLUDED_TOOLS.has(tool.name))
-                continue;
             const prefixedName = `${TOOL_PREFIX}${tool.name}`;
             const originalName = tool.name;
+            driverToolNames.add(prefixedName);
             pi.registerTool({
                 name: prefixedName,
                 label: prefixedName,
                 description: tool.description ?? '',
                 parameters: Type.Unsafe(tool.inputSchema),
-                async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
+                async execute(_toolCallId, params, signal, _onUpdate, ctx) {
+                    if (!(await confirmToolCall(originalName, params, ctx))) {
+                        return {
+                            content: [
+                                {
+                                    type: 'text',
+                                    text: `${originalName} requires interactive user confirmation.`,
+                                },
+                            ],
+                            details: undefined,
+                            isError: true,
+                        };
+                    }
                     try {
-                        await ensureConnected();
+                        await ensureConnected(signal, originalName !== 'check_permissions');
+                        const result = await client.callTool(originalName, params, signal);
+                        if (result.isError) {
+                            const errorText = result.content?.map((c) => c.text ?? '').join('') ?? '';
+                            const friendlyError = formatToolError(originalName, errorText, params);
+                            if (friendlyError) {
+                                return {
+                                    content: [{ type: 'text', text: friendlyError }],
+                                    details: undefined,
+                                    isError: true,
+                                };
+                            }
+                        }
+                        return toPiToolResult(result);
                     }
                     catch (connErr) {
+                        if (signal?.aborted)
+                            throw connErr;
                         const msg = connErr instanceof Error ? connErr.message : String(connErr);
                         ctx.ui.notify(`pi-computer-use: cannot connect to cua-driver — ${msg}`, 'warning');
                         return {
@@ -422,36 +355,61 @@ export default function computerUseExtension(pi) {
                             isError: true,
                         };
                     }
-                    const result = await client.callTool(originalName, params);
-                    if (result.isError) {
-                        const errorText = result.content?.map((c) => c.text ?? '').join('') ?? '';
-                        const friendlyError = formatToolError(originalName, errorText, params);
-                        if (friendlyError) {
-                            return {
-                                content: [{ type: 'text', text: friendlyError }],
-                                details: undefined,
-                                isError: true,
-                            };
-                        }
-                    }
-                    const content = [];
-                    if (result.content) {
-                        for (const item of result.content) {
-                            if (item.type === 'text' && item.text) {
-                                content.push({ type: 'text', text: item.text });
-                            }
-                        }
-                    }
-                    if (content.length === 0) {
-                        content.push({ type: 'text', text: 'Action executed.' });
-                    }
-                    return result.isError
-                        ? { content, details: undefined, isError: true }
-                        : { content, details: undefined };
                 },
             });
         }
     }
+    async function discoverAndRegisterTools(signal) {
+        await ensureConnected(signal);
+        const liveTools = await client.listAllTools(signal);
+        registerTools(liveTools);
+        return liveTools.length;
+    }
+    function registerRecoverySurface() {
+        pi.registerTool({
+            name: `${TOOL_PREFIX}connect`,
+            label: `${TOOL_PREFIX}connect`,
+            description: 'Retry Cua Driver startup and register the exact live tool contract for this platform.',
+            parameters: Type.Object({}),
+            async execute(_toolCallId, _params, signal) {
+                try {
+                    const count = await discoverAndRegisterTools(signal);
+                    applyToolProfile();
+                    return {
+                        content: [
+                            {
+                                type: 'text',
+                                text: `Cua Driver connected; registered ${count} platform tools. Continue with the requested computer_use tool.`,
+                            },
+                        ],
+                        details: { registered_tools: count },
+                    };
+                }
+                catch (error) {
+                    if (signal?.aborted)
+                        throw error;
+                    return {
+                        content: [{ type: 'text', text: 'Cua Driver is still unavailable.' }],
+                        details: undefined,
+                        isError: true,
+                    };
+                }
+            },
+        });
+        pi.registerCommand('computer-use-connect', {
+            description: 'Retry Cua Driver startup and discover platform tools',
+            async handler(_args, ctx) {
+                try {
+                    const count = await discoverAndRegisterTools(ctx.signal);
+                    applyToolProfile();
+                    ctx.ui.notify(`pi-computer-use: registered ${count} platform tools.`, 'info');
+                }
+                catch {
+                    ctx.ui.notify('pi-computer-use: Cua Driver is still unavailable.', 'warning');
+                }
+            },
+        });
+    }
     function registerVisionTool() {
         if (!config?.visionModel)
             return;
@@ -459,17 +417,23 @@ export default function computerUseExtension(pi) {
         pi.registerTool({
             name: `${TOOL_PREFIX}analyze_screenshot`,
             label: `${TOOL_PREFIX}analyze_screenshot`,
-            description: 'Take a screenshot and analyze it visually using a vision model. Use when you need to identify elements by visual attributes (color, layout, position) or need precise pixel coordinates.',
+            description: 'Capture a window through get_window_state and analyze its image using the configured vision model. Pass pid and window_id from list_windows.',
             parameters: Type.Object({
+                pid: Type.Integer({ description: 'Target process ID from list_windows.' }),
+                window_id: Type.Integer({
+                    description: 'Required CGWindowID / kCGWindowNumber to capture.',
+                }),
                 instruction: Type.Optional(Type.String({
                     description: 'What to identify or analyze visually (e.g., "Find the coordinates of the blue submit button").',
                 })),
             }),
-            async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
+            async execute(_toolCallId, params, signal, _onUpdate, ctx) {
                 try {
-                    await ensureConnected();
+                    await ensureConnected(signal);
                 }
                 catch (connErr) {
+                    if (signal?.aborted)
+                        throw connErr;
                     const msg = connErr instanceof Error ? connErr.message : String(connErr);
                     ctx.ui.notify(`pi-computer-use: cannot connect to cua-driver — ${msg}`, 'warning');
                     return {
@@ -483,11 +447,46 @@ export default function computerUseExtension(pi) {
                         isError: true,
                     };
                 }
-                const screenshotResult = await client.callTool('screenshot', {});
+                const screenshotArgs = {
+                    pid: params.pid,
+                    window_id: params.window_id,
+                    include_screenshot: true,
+                    max_elements: 1,
+                };
+                let screenshotResult;
+                try {
+                    screenshotResult = await client.callTool('get_window_state', screenshotArgs, signal);
+                }
+                catch (error) {
+                    if (signal?.aborted)
+                        throw error;
+                    const message = error instanceof Error ? error.message : String(error);
+                    return {
+                        content: [
+                            { type: 'text', text: `Failed to capture window state: ${message}` },
+                        ],
+                        details: undefined,
+                        isError: true,
+                    };
+                }
                 const imageContent = screenshotResult.content?.find((c) => c.type === 'image' && c.data);
+                if (process.env.DEBUG?.includes('pi-computer-use')) {
+                    console.error('[pi-computer-use analyze_screenshot] screenshot result', JSON.stringify({
+                        window_id: params.window_id,
+                        isError: screenshotResult.isError,
+                        contentTypes: screenshotResult.content?.map((c) => c.type),
+                        imageDataLength: imageContent?.data?.length,
+                        imageMimeType: imageContent?.mimeType,
+                    }, null, 2));
+                }
                 if (!imageContent?.data) {
+                    const errorText = screenshotResult.content
+                        ?.filter((c) => c.type === 'text' && c.text)
+                        .map((c) => c.text)
+                        .join('\n') || 'Failed to capture screenshot.';
+                    const formatted = formatToolError('get_window_state', errorText, params);
                     return {
-                        content: [{ type: 'text', text: 'Failed to capture screenshot.' }],
+                        content: [{ type: 'text', text: formatted ?? errorText }],
                         details: undefined,
                         isError: true,
                     };
@@ -495,7 +494,12 @@ export default function computerUseExtension(pi) {
                 const callVision = createPiVisionCaller(visionConfig, ctx);
                 const instruction = params.instruction ??
                     'Describe the full screen: identify all visible windows, UI elements, buttons, text fields, and their positions.';
-                const analysis = await callVision(instruction, imageContent.data, imageContent.mimeType ?? 'image/png');
+                const analysis = await callVision(instruction, imageContent.data, imageContent.mimeType ?? 'image/png', signal);
+                if (process.env.DEBUG?.includes('pi-computer-use')) {
+                    console.error('[pi-computer-use analyze_screenshot] vision analysis', JSON.stringify({
+                        analysisLength: analysis.length,
+                    }, null, 2));
+                }
                 return {
                     content: [{ type: 'text', text: analysis }],
                     details: undefined,
@@ -503,45 +507,61 @@ export default function computerUseExtension(pi) {
             },
         });
     }
+    registerToolGroupSurface();
     pi.on('session_start', async (_event, ctx) => {
-        config = resolveConfig(loadConfigFromFile({ cwd: ctx.cwd }));
-        client = new CuaDriverClient(config);
-        connected = false;
-        let upstreamTools;
+        config = resolveConfig(loadConfigFromFile({
+            cwd: ctx.cwd,
+            projectTrusted: isProjectTrusted(ctx),
+        }));
+        client = undefined;
+        sessionAbortController = new AbortController();
+        macPermissionPromise = undefined;
+        approvedLaunchApprovalKeys.clear();
+        driverToolNames.clear();
+        if (process.platform === 'darwin') {
+            registerTools(toolManifest.tools);
+            registerVisionTool();
+            applyToolProfile();
+            return;
+        }
+        let connectedAtStartup = false;
         try {
-            await client.connect();
-            connected = true;
-            upstreamTools = await client.listAllTools();
+            await discoverAndRegisterTools(ctx.signal);
+            connectedAtStartup = true;
         }
-        catch (err) {
-            const msg = err instanceof Error ? err.message : String(err);
-            ctx.ui.notify(`pi-computer-use: cua-driver failed to start — ${msg}. Tools registered but may fail until permissions are granted.`, 'warning');
+        catch (error) {
+            registerRecoverySurface();
+            const errorName = error instanceof Error ? error.name : 'UnknownError';
+            console.error(`[pi-computer-use] startup discovery failed (${errorName})`);
+            ctx.ui.notify('pi-computer-use: driver unavailable; use computer_use_connect to retry discovery.', 'warning');
         }
-        registerTools(upstreamTools ?? FALLBACK_TOOLS);
         registerVisionTool();
-        if (connected) {
-            try {
-                const permResult = await client.callTool('check_permissions', {});
-                const structured = permResult.structuredContent;
-                if (structured) {
-                    if (!structured.accessibility) {
-                        ctx.ui.notify(`pi-computer-use: ${accessibilityHint()}`, 'warning');
-                    }
-                    if (!structured.screen_recording) {
-                        ctx.ui.notify(`pi-computer-use: ${screenRecordingHint()}`, 'warning');
-                    }
-                }
-            }
-            catch {
-                // permission check is best-effort
+        applyToolProfile();
+        if (!connectedAtStartup)
+            return;
+        try {
+            const permissions = await client.callTool('check_permissions', { prompt: false }, ctx.signal);
+            const status = permissions.structuredContent;
+            if (status?.accessibility === false || status?.screen_recording === false) {
+                ctx.ui.notify(`pi-computer-use: ${permissionHint()}`, 'warning');
             }
         }
+        catch (error) {
+            const errorName = error instanceof Error ? error.name : 'UnknownError';
+            console.error(`[pi-computer-use] startup probe failed (${errorName})`);
+        }
     });
     pi.on('session_shutdown', async () => {
-        if (connected && client) {
-            await client.close();
-            connected = false;
-        }
+        const closingClient = client;
+        const sessionAbort = sessionAbortController;
+        const pendingPermission = macPermissionPromise;
+        client = undefined;
+        config = undefined;
+        sessionAbortController = undefined;
+        macPermissionPromise = undefined;
+        approvedLaunchApprovalKeys.clear();
+        sessionAbort?.abort(new Error('pi-computer-use: session shut down'));
+        await Promise.allSettled([pendingPermission, closingClient?.close()]);
     });
 }
 function formatToolError(toolName, errorText, params) {
@@ -551,7 +571,7 @@ function formatToolError(toolName, errorText, params) {
     if (errorText.includes('sc_not_granted')) {
         return screenRecordingHint();
     }
-    if (toolName === 'screenshot' || errorText.includes('screencapture failed')) {
+    if (toolName === 'get_window_state' || errorText.includes('screencapture failed')) {
         const windowId = params.window_id;
         if (windowId !== undefined) {
             if (errorText.includes('screencapture failed')) {
@@ -560,7 +580,7 @@ function formatToolError(toolName, errorText, params) {
                     `1. ${screenRecordingHint()}`,
                     '2. The window_id is stale — the window may have been closed or recreated (e.g. after navigation in Electron apps). Re-fetch window list to get current IDs.',
                     '3. The window is minimized or not yet rendered.',
-                    `Try capturing without window_id (full screen) as a fallback, or verify the window still exists.`,
+                    'Re-run list_windows and retry get_window_state with the current window_id.',
                 ].join('\n');
             }
             if (errorText.includes('empty output')) {