@celljs/ta-computer-use 0.0.0-stage → 4.8.0

package/src/extension.spec.ts

@@ -0,0 +1,279 @@
+import { describe, expect, it } from 'bun:test';
+import { CuaDriver, type CuaDriverProgram } from './cua-driver';
+import type { CuaDriverAccess, CuaDriverReadiness } from './cua-driver-access';
+import type { CuaDriverDaemon } from './cua-driver-daemon';
+import { activate, type ComputerUseOptions, type CuaDriverSetup } from './extension';
+import type { ExtensionAPI } from './pi-types';
+
+type Handler = (event: unknown, ctx: unknown) => unknown;
+
+/** The part of pi's extension API the extension uses; a registered server "connects" after `connectMs`. */
+class FakePi {
+    readonly handlers = new Map<string, Handler>();
+    readonly servers: { name: string; config: Record<string, unknown> }[] = [];
+    readonly tools: { name: string; reachesTools?: (input: Record<string, unknown>) => 'all' | readonly string[] }[] = [
+        {
+            name: 'codemode',
+            reachesTools: input => (/searchTools|describeNamespace/.test(String(input.code)) ? 'all' : (String(input.code).match(/\w+/g) ?? [])),
+        },
+        { name: 'bash' },
+    ];
+    readonly commands = new Map<string, { handler: (args: string, ctx: unknown) => Promise<void> }>();
+    readonly notes: string[] = [];
+    readonly events: string[] = [];
+    connectMs = 0;
+
+    get api(): ExtensionAPI {
+        return {
+            on: (event: string, handler: Handler) => this.handlers.set(event, handler),
+            getAllTools: () => this.tools,
+            getMcpServers: () => this.servers,
+            registerMcpServer: (name: string, config: Record<string, unknown>) => {
+                this.events.push('register');
+                this.servers.push({ name, config });
+                setTimeout(() => this.tools.push({ name: `${CuaDriver.NAMESPACE}__click` }), this.connectMs);
+            },
+            registerCommand: (name: string, command: { handler: (args: string, ctx: unknown) => Promise<void> }) => this.commands.set(name, command),
+        } as unknown as ExtensionAPI;
+    }
+
+    get ctx(): unknown {
+        return { ui: { notify: (text: string) => this.notes.push(text) } };
+    }
+
+    prompt(): Record<string, string> {
+        const sections: Record<string, string> = {};
+        this.handlers.get('before_agent_start')!({ systemPromptOptions: { sections } }, {});
+        return sections;
+    }
+
+    call(toolName: string, input: Record<string, unknown>, signal?: AbortSignal): Promise<unknown> {
+        return Promise.resolve(this.handlers.get('tool_call')!({ toolName, input }, { ...(this.ctx as object), signal }));
+    }
+}
+
+const CLICK = { code: 'await tools.mcp__cua_driver__click({ x: 1 })' };
+const READY: CuaDriverReadiness = { ready: true, notes: [], known: true };
+const OWN: CuaDriverProgram = { path: '/bin/cua-driver', bundled: false };
+
+type FakeSetup = CuaDriverSetup & { checks: number; ensured: number };
+
+/** A setup whose check answers `answers` in turn (the last one again), and whose daemon, if any, counts its starts. */
+function fakeSetup(pi: FakePi, program: CuaDriverProgram, answers: CuaDriverReadiness[], withDaemon = false): FakeSetup {
+    const setup = {
+        program,
+        args: withDaemon ? ['mcp', '--embedded', '--socket', '/c/s.sock'] : ['mcp'],
+        ready: false,
+        checks: 0,
+        ensured: 0,
+    } as unknown as FakeSetup;
+    setup.access = {
+        check: async (notify: (message: string) => void) => {
+            pi.events.push('check');
+            const answer = answers[Math.min(setup.checks++, answers.length - 1)];
+            if (!answer.ready) {
+                notify('asked');
+            }
+            return answer;
+        },
+    } as unknown as CuaDriverAccess;
+    if (withDaemon) {
+        setup.daemon = {
+            socket: '/c/s.sock',
+            ensure: async () => {
+                pi.events.push('ensure');
+                setup.ensured++;
+            },
+            reap: async () => undefined,
+        } as unknown as CuaDriverDaemon;
+    }
+    return setup;
+}
+
+function options(
+    locate: () => CuaDriverProgram | undefined,
+    setup: (program: CuaDriverProgram) => CuaDriverSetup,
+    platform: NodeJS.Platform = 'linux',
+): ComputerUseOptions {
+    return { platform, locate, setup };
+}
+
+describe('the computer-use extension', () => {
+    it('without Cua Driver, adds nothing to the prompt and registers nothing', async () => {
+        const pi = new FakePi();
+        activate(
+            pi.api,
+            options(
+                () => undefined,
+                () => {
+                    throw new Error('no setup');
+                },
+            ),
+        );
+        expect(pi.prompt()).toEqual({});
+        await pi.call('codemode', CLICK);
+        expect(pi.servers).toEqual([]);
+    });
+
+    it('names the driver in the prompt, and starts it only for a call that names its tools', async () => {
+        const pi = new FakePi();
+        let looked = 0;
+        const setup = fakeSetup(pi, OWN, [READY]);
+        activate(
+            pi.api,
+            options(
+                () => {
+                    looked++;
+                    return OWN;
+                },
+                () => setup,
+            ),
+        );
+        expect(pi.prompt()[CuaDriver.SECTION]).toBe(CuaDriver.prompt('linux'));
+        pi.prompt();
+        expect(looked).toBe(1);
+        await pi.call('bash', { command: 'ls' });
+        await pi.call('codemode', { code: 'await tools.bash({ command: "ls" })' });
+        await pi.call('codemode', { code: 'searchTools("click")' });
+        expect(pi.servers).toEqual([]);
+        pi.connectMs = 40;
+        expect(await pi.call('codemode', CLICK)).toBeUndefined();
+        expect(pi.servers).toEqual([
+            { name: CuaDriver.SERVER, config: { command: '/bin/cua-driver', args: ['mcp'], env: undefined, description: CuaDriver.DESCRIPTION } },
+        ]);
+        expect(pi.tools.some(tool => tool.name === 'mcp__cua_driver__click')).toBe(true);
+        await pi.call('codemode', { code: 'describeNamespace("mcp__cua_driver")' });
+        expect(pi.servers.length).toBe(1);
+        // The user's own install starts its daemon when it connects: connected first, then checked, once.
+        expect(pi.events).toEqual(['register', 'check']);
+    });
+
+    it('registers and checks once for calls at the same time', async () => {
+        const pi = new FakePi();
+        const setup = fakeSetup(pi, OWN, [READY]);
+        activate(
+            pi.api,
+            options(
+                () => OWN,
+                () => setup,
+            ),
+        );
+        pi.prompt();
+        pi.connectMs = 20;
+        await Promise.all([pi.call('codemode', CLICK), pi.call('codemode', CLICK)]);
+        expect(pi.servers.length).toBe(1);
+        expect(setup.checks).toBe(1);
+    });
+
+    it('looks again on each prompt until the driver is installed', () => {
+        const pi = new FakePi();
+        let installed: CuaDriverProgram | undefined;
+        activate(
+            pi.api,
+            options(
+                () => installed,
+                program => fakeSetup(pi, program, [READY]),
+            ),
+        );
+        expect(pi.prompt()).toEqual({});
+        installed = OWN;
+        expect(pi.prompt()[CuaDriver.SECTION]).toBeDefined();
+    });
+
+    it('stops waiting for the server when the call is stopped', async () => {
+        const pi = new FakePi();
+        activate(
+            pi.api,
+            options(
+                () => OWN,
+                program => fakeSetup(pi, program, [READY]),
+            ),
+        );
+        pi.prompt();
+        pi.connectMs = 10_000;
+        const stop = new AbortController();
+        const started = Date.now();
+        const call = pi.call('codemode', CLICK, stop.signal);
+        setTimeout(() => stop.abort(), 30);
+        await call;
+        expect(Date.now() - started).toBeLessThan(1000);
+        expect(pi.servers.length).toBe(1);
+    });
+
+    it('stops a call with what the system lacks, checks again on the next, and passes the notes on once', async () => {
+        const pi = new FakePi();
+        const setup = fakeSetup(pi, OWN, [
+            { ready: false, reason: 'no AT-SPI' },
+            { ready: true, notes: ['AT-SPI: slow'], known: true },
+        ]);
+        activate(
+            pi.api,
+            options(
+                () => OWN,
+                () => setup,
+            ),
+        );
+        pi.prompt();
+        expect(await pi.call('codemode', CLICK)).toEqual({ block: true, reason: 'no AT-SPI' });
+        expect(await pi.call('codemode', CLICK)).toBeUndefined();
+        await pi.call('codemode', CLICK);
+        expect(pi.notes).toEqual(['asked', 'AT-SPI: slow']);
+        expect(setup.checks).toBe(2);
+    });
+
+    it('asks again when the system could not tell', async () => {
+        const pi = new FakePi();
+        const setup = fakeSetup(pi, OWN, [{ ready: true, notes: [], known: false }, READY]);
+        activate(
+            pi.api,
+            options(
+                () => OWN,
+                () => setup,
+            ),
+        );
+        pi.prompt();
+        await pi.call('codemode', CLICK);
+        await pi.call('codemode', CLICK);
+        await pi.call('codemode', CLICK);
+        expect(setup.checks).toBe(2);
+    });
+
+    it("with Ta Desktop's daemon, starts it and checks before the server connects, and keeps it running", async () => {
+        const pi = new FakePi();
+        const program = { path: '/T/CuaDriver.app/Contents/MacOS/cua-driver', bundled: true };
+        const setup = fakeSetup(pi, program, [READY], true);
+        activate(
+            pi.api,
+            options(
+                () => program,
+                () => setup,
+                'darwin',
+            ),
+        );
+        pi.prompt();
+        await pi.call('codemode', { code: 'tools.bash({})' });
+        expect(setup.ensured).toBe(0);
+        await pi.call('codemode', CLICK);
+        await pi.call('codemode', CLICK);
+        expect(pi.events).toEqual(['ensure', 'check', 'register', 'ensure']);
+        expect(pi.servers[0].config.args).toEqual(['mcp', '--embedded', '--socket', '/c/s.sock']);
+    });
+
+    it('/computer-use says where the driver is or how to install it, and what the system lacks', async () => {
+        const pi = new FakePi();
+        let installed: CuaDriverProgram | undefined;
+        const setup = fakeSetup(pi, OWN, [{ ready: false, reason: 'no display' }]);
+        activate(
+            pi.api,
+            options(
+                () => installed,
+                () => setup,
+            ),
+        );
+        await pi.commands.get('computer-use')!.handler('', pi.ctx);
+        expect(pi.notes[0]).toContain('cua.ai/driver/install');
+        installed = OWN;
+        await pi.commands.get('computer-use')!.handler('', pi.ctx);
+        expect(pi.notes.slice(1)).toEqual([CuaDriver.status(OWN, 'linux'), 'asked', 'no display']);
+    });
+});

package/src/extension.ts

@@ -0,0 +1,195 @@
+import { execFile } from 'node:child_process';
+import { realpathSync } from 'node:fs';
+import { join } from 'node:path';
+import { CuaDriver, type CuaDriverProgram } from './cua-driver';
+import { CuaDriverAccess } from './cua-driver-access';
+import { CuaDriverDaemon } from './cua-driver-daemon';
+import type { ExtensionAPI } from './pi-types';
+
+/** How long a call waits the first time for Cua Driver's tools; it connects in tens of milliseconds. */
+const CONNECT_WAIT_MS = 5000;
+const POLL_MS = 15;
+
+/** A program of Cua Driver and what goes with it, shared by the sessions of a process. */
+export interface CuaDriverSetup {
+    program: CuaDriverProgram;
+    /** The daemon Ta runs, for the driver of Ta Desktop on macOS. */
+    daemon?: CuaDriverDaemon;
+    access: CuaDriverAccess;
+    /** The arguments and environment of its MCP server. */
+    args: string[];
+    env?: Record<string, string>;
+    /** Checked, once for the process. */
+    ready: boolean;
+    checking?: Promise<Awaited<ReturnType<CuaDriverAccess['check']>>>;
+}
+
+/** What the extension depends on, so that tests can give their own. */
+export interface ComputerUseOptions {
+    locate(): CuaDriverProgram | undefined;
+    platform: NodeJS.Platform;
+    setup(program: CuaDriverProgram, platform: NodeJS.Platform): CuaDriverSetup;
+}
+
+const setups = new Map<string, CuaDriverSetup>();
+
+/** One per program and process: the sessions of `ta serve` share its daemon and its check. */
+export function setupOf(found: CuaDriverProgram, platform: NodeJS.Platform): CuaDriverSetup {
+    // LaunchServices finds an app by its exact path: one with `..` or a link in it is another app to it.
+    const program = { ...found, path: realPath(found.path) };
+    let setup = setups.get(program.path);
+    if (!setup) {
+        const env = program.bundled ? { ...CuaDriver.BUNDLED_ENV_VARS } : undefined;
+        const app = program.bundled && platform === 'darwin' ? CuaDriver.appOf(program.path) : undefined;
+        const daemon = app ? new CuaDriverDaemon(app, { env }) : undefined;
+        const access = new CuaDriverAccess(daemon ? join(daemon.app, 'Contents', 'MacOS', 'cua-driver') : program.path, {
+            platform,
+            socket: daemon?.socket,
+            env,
+            beforeLook: daemon ? () => daemon.ensure() : undefined,
+            openSettings: pane => execFile('open', [`x-apple.systempreferences:com.apple.preference.security?${pane}`], () => undefined),
+        });
+        const args = daemon ? ['mcp', '--embedded', '--socket', daemon.socket] : ['mcp'];
+        setup = { program, daemon, access, args, env, ready: false };
+        setups.set(program.path, setup);
+    }
+    return setup;
+}
+
+function realPath(path: string): string {
+    try {
+        return realpathSync(path);
+    } catch {
+        return path;
+    }
+}
+
+const DEFAULTS: ComputerUseOptions = { locate: () => CuaDriver.locate(), platform: process.platform, setup: setupOf };
+
+/**
+ * Lets the agent use the desktop apps of this computer through Cua Driver, at no cost until it does.
+ * Registered up front, an MCP server starts in every session: its process (about 10 MB) and pi's MCP
+ * client. So the system prompt only names it, and the server is registered when a call first names its
+ * tools; pi connects it before the call runs. Each prompt looks for the program (a few `stat`s, nothing
+ * is run) until it is found, so a driver installed meanwhile is used without a restart.
+ *
+ * The first call also checks what the system needs (CuaDriverAccess): on macOS it waits while the user
+ * gives the permissions it asked for, elsewhere it says what is missing, so that the agent's call goes on
+ * or fails with what to do, not with an error of the driver.
+ */
+export function activate(pi: ExtensionAPI, options: ComputerUseOptions = DEFAULTS): void {
+    let setup: CuaDriverSetup | undefined;
+    /** Registered, and waited for once: later calls are pi's to wait for, as for any server it knows. */
+    let connected = false;
+    let connecting: Promise<void> | undefined;
+    const prefix = `${CuaDriver.NAMESPACE}__`;
+    const ready = () => pi.getAllTools().some(tool => tool.name.startsWith(prefix));
+
+    const found = (): CuaDriverSetup | undefined => {
+        if (!setup) {
+            const program = options.locate();
+            setup = program ? options.setup(program, options.platform) : undefined;
+        }
+        return setup;
+    };
+
+    const connect = async (driver: CuaDriverSetup, signal: AbortSignal | undefined): Promise<void> => {
+        if (!pi.getMcpServers().some(server => server.name === CuaDriver.SERVER)) {
+            pi.registerMcpServer(CuaDriver.SERVER, { command: driver.program.path, args: driver.args, env: driver.env, description: CuaDriver.DESCRIPTION });
+        }
+        // pi's MCP extension waits for a server it already knows; this one may come after this handler.
+        const deadline = Date.now() + CONNECT_WAIT_MS;
+        while (!ready() && Date.now() < deadline && !signal?.aborted) {
+            await new Promise(resolve => setTimeout(resolve, POLL_MS));
+        }
+        // A call stopped meanwhile leaves the wait to the next one.
+        connected = !signal?.aborted;
+    };
+
+    /** Once for the process; the calls at the same time share it. */
+    const check = async (driver: CuaDriverSetup, notify: (message: string) => void, signal: AbortSignal | undefined): Promise<string | undefined> => {
+        driver.checking ??= driver.access.check(notify, signal).finally(() => {
+            driver.checking = undefined;
+        });
+        const result = await driver.checking;
+        if (!result.ready) {
+            return result.reason;
+        }
+        if (result.known && !driver.ready) {
+            driver.ready = true;
+            for (const note of result.notes) {
+                notify(note);
+            }
+            void driver.daemon?.reap();
+        }
+        return undefined;
+    };
+
+    pi.on('before_agent_start', event => {
+        const { sections } = event.systemPromptOptions;
+        if (found()) {
+            sections[CuaDriver.SECTION] = CuaDriver.prompt(options.platform);
+        } else {
+            delete sections[CuaDriver.SECTION];
+        }
+    });
+
+    pi.on('tool_call', async (event, ctx) => {
+        const driver = setup;
+        if (!driver) {
+            return undefined;
+        }
+        if (driver.ready && connected && !driver.daemon) {
+            return undefined;
+        }
+        const tool = pi.getAllTools().find(candidate => candidate.name === event.toolName);
+        if (!CuaDriver.reaches(tool?.reachesTools?.(event.input), event.input)) {
+            return undefined;
+        }
+        const notify = (message: string) => ctx.ui.notify(message, 'warning');
+        try {
+            // Every call that uses it, in case the daemon stopped: a check costs one connection.
+            await driver.daemon?.ensure();
+            // The daemon of Ta Desktop answers before the server connects: its relay needs the permissions.
+            // The user's own install starts its daemon when it connects, so it is asked afterwards.
+            const reason = driver.daemon && !driver.ready ? await check(driver, notify, ctx.signal) : undefined;
+            if (reason) {
+                return { block: true, reason };
+            }
+            if (!connected) {
+                connecting ??= connect(driver, ctx.signal).finally(() => {
+                    connecting = undefined;
+                });
+                await connecting;
+            }
+            const after = !driver.daemon && !driver.ready ? await check(driver, notify, ctx.signal) : undefined;
+            return after ? { block: true, reason: after } : undefined;
+        } catch (error) {
+            return { block: true, reason: (error as Error).message };
+        }
+    });
+
+    pi.registerCommand('computer-use', {
+        description: 'Whether Ta can use the desktop apps of this computer (Cua Driver), and what the system still needs',
+        handler: async (_args, ctx) => {
+            const program = options.locate();
+            ctx.ui.notify(CuaDriver.status(program, options.platform), program ? 'info' : 'warning');
+            if (!program) {
+                return;
+            }
+            const driver = options.setup(program, options.platform);
+            await driver.daemon?.ensure();
+            const result = await driver.access.check(message => ctx.ui.notify(message, 'warning'), ctx.signal);
+            if (!result.ready) {
+                ctx.ui.notify(result.reason, 'warning');
+            } else if (result.known) {
+                driver.ready = true;
+                ctx.ui.notify(['Cua Driver has what it needs.', ...result.notes].join('\n'), result.notes.length > 0 ? 'warning' : 'info');
+            }
+        },
+    });
+}
+
+export default function computerUse(pi: ExtensionAPI): void {
+    activate(pi);
+}

package/src/pi-types.ts

@@ -0,0 +1,2 @@
+/** pi's types, which are ES modules, for this package's CommonJS build; only types, so the bundle imports nothing of pi. */
+export type ExtensionAPI = import('@earendil-works/pi-coding-agent', { with: { 'resolution-mode': 'import' }}).ExtensionAPI;