flynn/src/companion/platformClients.integration.test.ts

import { describe, it, expect, beforeAll, afterAll, vi } from 'vitest';
import { resolve } from 'path';
import { createServer } from 'net';
import type { GatewayServerConfig } from '../gateway/server.js';
import { GatewayServer } from '../gateway/server.js';
import { CompanionRuntimeClient } from './runtimeClient.js';
import {
  AndroidCompanionClient,
  IOSCompanionClient,
  MacOSCompanionClient,
} from './platformClients.js';

async function canListenOnLocalhost(): Promise<boolean> {
  return new Promise((resolvePromise) => {
    const s = createServer();
    s.once('error', () => resolvePromise(false));
    s.listen(0, '127.0.0.1', () => {
      s.close(() => resolvePromise(true));
    });
  });
}

const mockSession = {
  id: 'test',
  addMessage: vi.fn(),
  getHistory: vi.fn(() => []),
  clear: vi.fn(),
  setHistory: vi.fn(),
  replaceHistory: vi.fn(),
};

const mockSessionManager = {
  getSession: vi.fn(() => mockSession),
  getSessionConfig: vi.fn(() => undefined),
  listSessions: vi.fn(() => ['ws:test']),
  transferSession: vi.fn(),
  closeSession: vi.fn(),
};

const mockModelClient = {
  chat: vi.fn(async () => ({
    content: 'Hello from Flynn!',
    stopReason: 'end_turn',
    usage: { inputTokens: 10, outputTokens: 5 },
  })),
};

const mockToolRegistry = {
  register: vi.fn(),
  get: vi.fn(),
  list: vi.fn(() => []),
  filteredList: vi.fn(() => []),
  toAnthropicFormat: vi.fn(() => []),
  toOpenAIFormat: vi.fn(() => []),
  filteredToAnthropicFormat: vi.fn(() => []),
  filteredToOpenAIFormat: vi.fn(() => []),
};

const mockToolExecutor = {
  execute: vi.fn(async () => ({ success: true, output: 'ok' })),
};

const TEST_PORT = 18912;
const TEST_TOKEN = 'platform-clients-token';

let LISTEN_ALLOWED = true;
let server: GatewayServer;

beforeAll(async () => {
  LISTEN_ALLOWED = await canListenOnLocalhost();
  if (!LISTEN_ALLOWED) {
    return;
  }

  server = new GatewayServer({
    port: TEST_PORT,
    sessionManager: mockSessionManager as unknown as GatewayServerConfig['sessionManager'],
    modelClient: mockModelClient,
    systemPrompt: 'Test prompt',
    toolRegistry: mockToolRegistry as unknown as GatewayServerConfig['toolRegistry'],
    toolExecutor: mockToolExecutor as unknown as GatewayServerConfig['toolExecutor'],
    version: '0.1.0-test',
    uiDir: resolve(import.meta.dirname, '../gateway/ui'),
    auth: { token: TEST_TOKEN },
    nodes: {
      enabled: true,
      allowedRoles: ['companion'],
      featureGates: { 'ui.canvas': true },
      locationEnabled: true,
      pushEnabled: true,
    },
  });

  await server.start();
});

afterAll(async () => {
  if (!LISTEN_ALLOWED) {
    return;
  }
  await server.stop();
});

function createRuntime(): CompanionRuntimeClient {
  return new CompanionRuntimeClient({
    url: `ws://127.0.0.1:${TEST_PORT}`,
    token: TEST_TOKEN,
  });
}

describe('platform clients integration', () => {
  it('platform event subscription helpers track subscription count lifecycle', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new MacOSCompanionClient({ runtime, nodeId: 'macos-events-e2e' });

    expect(client.eventSubscriptionCount).toBe(0);
    const unsubscribeA = client.subscribeEvents(() => undefined);
    const unsubscribeB = client.subscribeEvent('agent.stream', () => undefined);
    expect(client.eventSubscriptionCount).toBe(2);

    unsubscribeA();
    expect(client.eventSubscriptionCount).toBe(1);

    client.clearEventSubscriptions();
    expect(client.eventSubscriptionCount).toBe(0);

    unsubscribeB();
    expect(client.eventSubscriptionCount).toBe(0);
  });

  it('platform pendingEventWaitCount tracks waitForAnyEvent lifecycle', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new IOSCompanionClient({ runtime, nodeId: 'ios-wait-count-e2e' });

    expect(client.pendingEventWaitCount).toBe(0);
    const awaited = expect(
      client.waitForAnyEvent(['agent.stream'], { timeoutMs: 10_000 }),
    ).rejects.toThrow('Event subscriptions cleared');
    expect(client.pendingEventWaitCount).toBe(1);

    client.clearEventSubscriptions();
    await awaited;
    expect(client.pendingEventWaitCount).toBe(0);
  });

  it('platform waitForIdle resolves after pending waiters are cleared', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new IOSCompanionClient({ runtime, nodeId: 'ios-wait-idle-e2e' });

    const pending = client.waitForAnyEvent(['agent.stream'], { timeoutMs: 10_000 }).catch(() => undefined);
    expect(client.hasPendingWork).toBe(true);

    const idle = client.waitForIdle({ timeoutMs: 1_000, pollIntervalMs: 5 });
    setTimeout(() => {
      client.clearEventSubscriptions();
    }, 20);

    await expect(idle).resolves.toBeUndefined();
    await pending;
    expect(client.hasPendingWork).toBe(false);
  });

  it('platform getPendingWorkSnapshot reflects waiter lifecycle', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new IOSCompanionClient({ runtime, nodeId: 'ios-snapshot-e2e' });

    expect(client.getPendingWorkSnapshot()).toEqual({
      pendingRequestCount: 0,
      pendingEventWaitCount: 0,
      hasPendingWork: false,
    });

    const pending = client.waitForAnyEvent(['agent.stream'], { timeoutMs: 10_000 }).catch(() => undefined);
    expect(client.getPendingWorkSnapshot()).toEqual({
      pendingRequestCount: 0,
      pendingEventWaitCount: 1,
      hasPendingWork: true,
    });

    client.clearEventSubscriptions();
    await pending;
    expect(client.getPendingWorkSnapshot()).toEqual({
      pendingRequestCount: 0,
      pendingEventWaitCount: 0,
      hasPendingWork: false,
    });
  });

  it('platform cancelPendingEventWaits cancels waits without clearing subscriptions', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new IOSCompanionClient({ runtime, nodeId: 'ios-cancel-waits-e2e' });

    const unsubscribe = client.subscribeEvents(() => undefined);
    expect(client.eventSubscriptionCount).toBe(1);

    const awaited = expect(
      client.waitForAnyEvent(['agent.stream'], { timeoutMs: 10_000 }),
    ).rejects.toThrow('manual cancel');
    expect(client.pendingEventWaitCount).toBe(1);

    const cancelled = client.cancelPendingEventWaits('manual cancel');
    await awaited;
    expect(cancelled).toBe(1);
    expect(client.pendingEventWaitCount).toBe(0);
    expect(client.eventSubscriptionCount).toBe(1);
    expect(client.cancelPendingEventWaits()).toBe(0);

    unsubscribe();
    expect(client.eventSubscriptionCount).toBe(0);
  });

  it('platform getEventSurfaceSnapshot reflects subscription and waiter lifecycle', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new IOSCompanionClient({ runtime, nodeId: 'ios-event-snapshot-e2e' });

    const unsubscribe = client.subscribeEvents(() => undefined);
    const pending = client.waitForAnyEvent(['agent.stream'], { timeoutMs: 10_000 }).catch(() => undefined);

    expect(client.getEventSurfaceSnapshot()).toEqual({
      knownEventNames: ['agent.stream', 'agent.typing', 'context_warning'],
      eventSubscriptionCount: 2,
      pendingEventWaitCount: 1,
    });

    client.clearEventSubscriptions();
    await pending;
    expect(client.getEventSurfaceSnapshot()).toEqual({
      knownEventNames: ['agent.stream', 'agent.typing', 'context_warning'],
      eventSubscriptionCount: 0,
      pendingEventWaitCount: 0,
    });

    unsubscribe();
  });

  it('platform getConnectionSnapshot reflects connection and activity state', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new IOSCompanionClient({ runtime, nodeId: 'ios-connection-snapshot-e2e' });

    expect(client.getConnectionSnapshot()).toEqual({
      connected: false,
      eventSubscriptionCount: 0,
      pendingRequestCount: 0,
      pendingEventWaitCount: 0,
      hasPendingWork: false,
      idle: true,
      lastDisconnectCode: undefined,
      lastDisconnectReason: undefined,
    });

    const unsubscribe = client.subscribeEvents(() => undefined);
    const pending = client.waitForAnyEvent(['agent.stream'], { timeoutMs: 10_000 }).catch(() => undefined);

    expect(client.getConnectionSnapshot()).toEqual({
      connected: false,
      eventSubscriptionCount: 2,
      pendingRequestCount: 0,
      pendingEventWaitCount: 1,
      hasPendingWork: true,
      idle: false,
      lastDisconnectCode: undefined,
      lastDisconnectReason: undefined,
    });

    client.clearEventSubscriptions();
    await pending;
    expect(client.getConnectionSnapshot()).toEqual({
      connected: false,
      eventSubscriptionCount: 0,
      pendingRequestCount: 0,
      pendingEventWaitCount: 0,
      hasPendingWork: false,
      idle: true,
      lastDisconnectCode: undefined,
      lastDisconnectReason: undefined,
    });

    unsubscribe();
  });

  it('platform clearEventSubscriptions returns cleared and cancelled counts', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new IOSCompanionClient({ runtime, nodeId: 'ios-clear-counts-e2e' });

    const unsubscribe = client.subscribeEvents(() => undefined);
    const pending = client.waitForAnyEvent(['agent.stream'], { timeoutMs: 10_000 }).catch(() => undefined);

    const cleared = client.clearEventSubscriptions();
    expect(cleared).toEqual({
      clearedSubscriptions: 2,
      cancelledWaits: 1,
    });

    await pending;
    expect(client.eventSubscriptionCount).toBe(0);
    unsubscribe();
  });

  it('platform connected reflects runtime connection lifecycle', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new IOSCompanionClient({ runtime, nodeId: 'ios-connected-e2e' });

    expect(client.connected).toBe(false);
    await client.connect();
    expect(client.connected).toBe(true);
    client.disconnect(4100, 'manual platform stop');
    expect(client.connected).toBe(false);
    expect(client.lastDisconnectCode).toBe(4100);
    expect(client.lastDisconnectReason).toBe('manual platform stop');
    expect(client.getConnectionSnapshot()).toEqual({
      connected: false,
      eventSubscriptionCount: 0,
      pendingRequestCount: 0,
      pendingEventWaitCount: 0,
      hasPendingWork: false,
      idle: true,
      lastDisconnectCode: 4100,
      lastDisconnectReason: 'manual platform stop',
    });
  });

  it('macOS companion wrapper registers and writes status with platform pinning', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new MacOSCompanionClient({ runtime, nodeId: 'macos-e2e' });
    await client.connect();

    try {
      const boot = await client.bootstrap();
      expect(boot.register.registered).toBe(true);
      expect(boot.capabilities.node.id).toBe('macos-e2e');
      const status = await client.publishHeartbeat({
        appVersion: '1.0.0',
        statusText: 'menu-bar-active',
        powerSource: 'ac',
      });

      expect(status.updated).toBe(true);
      expect(status.status.platform).toBe('macos');

      const nodes = await client.listNodes();
      const entry = nodes.nodes.find((n) => n.nodeId === 'macos-e2e');
      expect(entry?.status?.platform).toBe('macos');
    } finally {
      client.disconnect();
    }
  });

  it('iOS companion wrapper uses APNs push flow', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new IOSCompanionClient({ runtime, nodeId: 'ios-e2e' });
    await client.connect();

    try {
      await client.register();
      const push = await client.registerPushToken({
        token: 'd'.repeat(64),
        topic: 'dev.flynn.ios',
        environment: 'sandbox',
      });

      expect(push.updated).toBe(true);
      expect(push.push.provider).toBe('apns');

      const nodes = await client.listNodes();
      const entry = nodes.nodes.find((n) => n.nodeId === 'ios-e2e');
      expect(entry?.push?.provider).toBe('apns');
    } finally {
      client.disconnect();
    }
  });

  it('Android companion wrapper uses FCM push flow', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new AndroidCompanionClient({ runtime, nodeId: 'android-e2e' });
    await client.connect();

    try {
      await client.register();
      const push = await client.registerPushToken('e'.repeat(64));

      expect(push.updated).toBe(true);
      expect(push.push.provider).toBe('fcm');

      const nodes = await client.listNodes();
      const entry = nodes.nodes.find((n) => n.nodeId === 'android-e2e');
      expect(entry?.push?.provider).toBe('fcm');
    } finally {
      client.disconnect();
    }
  });

  it('macOS companion wrapper supports canvas artifact lifecycle', async () => {
    if (!LISTEN_ALLOWED) {
      return;
    }

    const runtime = createRuntime();
    const client = new MacOSCompanionClient({
      runtime,
      nodeId: 'macos-canvas-e2e',
      defaultSessionId: 'ws:platform-canvas-e2e',
    });
    await client.connect();

    try {
      await client.register();

      const put = await client.putCanvasArtifact({
        artifactId: 'mac-art-1',
        type: 'note',
        content: { text: 'hello from platform wrapper' },
      });
      expect(put.upserted).toBe(true);
      expect(put.artifact.id).toBe('mac-art-1');

      const list = await client.listCanvasArtifacts();
      expect(list.artifacts.some((artifact) => artifact.id === 'mac-art-1')).toBe(true);

      const get = await client.getCanvasArtifact({ artifactId: 'mac-art-1' });
      expect(get.artifact.id).toBe('mac-art-1');

      const del = await client.deleteCanvasArtifact({ artifactId: 'mac-art-1' });
      expect(del.deleted).toBe(true);

      const clear = await client.clearCanvasArtifacts();
      expect(clear.cleared).toBe(0);
    } finally {
      client.disconnect();
    }
  });
});