picogent/src/index.ts

import { loadConfig, hasAgentConfig, loadAgentConfig } from './config.js';
import type { AgentConfig } from './config.js';
import { logger } from './logger.js';
import { runAgent, initAgent } from './agent.js';
import { WebSocketTransport } from './transport/websocket.js';
import { HttpTransport } from './transport/http.js';
import { WsClientTransport } from './transport/ws-client.js';
import { TrackerClient } from './tracker/client.js';
import { TrackerRegistration } from './tracker/registration.js';
import { EventRouter } from './router.js';
import type { IncomingMessage } from './transport/types.js';

async function startAgentHttp(config: AgentConfig, client: TrackerClient): Promise<void> {
  const registration = new TrackerRegistration(config);
  const router = new EventRouter(config, client);
  const http = new HttpTransport(config.listenPort);

  http.onEvent((event) => router.handleEvent(event));

  await http.start();

  const callbackUrl = `http://localhost:${config.listenPort}/events`;
  registration.register(callbackUrl);
  registration.startHeartbeat();

  const shutdown = () => {
    logger.info('Shutting down agent (http)...');
    registration.stopHeartbeat();
    http.stop().then(() => {
      logger.info('Agent stopped');
      process.exit(0);
    });
  };

  process.on('SIGINT', shutdown);
  process.on('SIGTERM', shutdown);
}

async function startAgentWs(config: AgentConfig, client: TrackerClient): Promise<void> {
  const wsTransport = new WsClientTransport(config);
  const router = new EventRouter(config, client);

  wsTransport.onEvent((event) => router.handleEvent(event));

  await wsTransport.start();
  logger.info('Connected to tracker via WebSocket');

  const shutdown = () => {
    logger.info('Shutting down agent (ws)...');
    wsTransport.stop().then(() => {
      logger.info('Agent stopped');
      process.exit(0);
    });
  };

  process.on('SIGINT', shutdown);
  process.on('SIGTERM', shutdown);
}

async function startAgentMode(): Promise<void> {
  const config = loadAgentConfig();
  initAgent(config.provider, config.apiKey);

  logger.info({
    config: {
      name: config.name,
      slug: config.slug,
      trackerUrl: config.trackerUrl,
      transport: config.transport,
      listenPort: config.transport === 'http' ? config.listenPort : undefined,
      workDir: config.workDir,
      agentHome: config.agentHome,
      model: config.model,
    },
  }, 'Starting picogent in agent mode');

  const client = new TrackerClient(config.trackerUrl, config.token);

  if (config.transport === 'ws') {
    await startAgentWs(config, client);
  } else {
    await startAgentHttp(config, client);
  }
}

async function startWebSocketMode(): Promise<void> {
  const config = loadConfig();
  initAgent(config.provider, config.apiKey);
  const transport = new WebSocketTransport(config.port);

  transport.onMessage(async (msg: IncomingMessage & { _clientId?: string }) => {
    const clientId = msg._clientId || 'unknown';
    const workDir = msg.workDir || config.defaultWorkDir;

    logger.info({
      messageId: msg.id,
      clientId,
      workDir,
      sessionId: msg.sessionId || '(none)',
    }, 'Processing message');

    for await (const agentMsg of runAgent(msg.content, {
      workDir,
      sessionId: msg.sessionId,
      model: config.model,
      provider: config.provider,
      systemPrompt: msg.systemPrompt,
    })) {
      transport.send(clientId, {
        id: msg.id,
        type: agentMsg.type === 'session' ? 'text' : agentMsg.type === 'done' ? 'done' : agentMsg.type,
        content: agentMsg.content,
        sessionId: agentMsg.sessionId,
      });
    }
  });

  logger.info({
    config: { port: config.port, defaultWorkDir: config.defaultWorkDir, sessionDir: config.sessionDir, model: config.model },
  }, 'Starting picogent');
  await transport.start();

  const shutdown = () => {
    logger.info('Shutting down...');
    transport.stop().then(() => {
      logger.info('Stopped');
      process.exit(0);
    });
  };

  process.on('SIGINT', shutdown);
  process.on('SIGTERM', shutdown);
}

async function main(): Promise<void> {
  if (hasAgentConfig()) {
    await startAgentMode();
  } else {
    await startWebSocketMode();
  }
}

main().catch((err) => {
  logger.fatal({ err }, 'Failed to start');
  process.exit(1);
});