ktx/packages/cli/src/agent.ts

import { readFile } from 'node:fs/promises';
import type { KloCliIo } from './cli-runtime.js';
import {
  createKloAgentRuntime,
  parseAgentMaxRows,
  readAgentJsonFile,
  writeAgentJson,
  writeAgentJsonError,
  type KloAgentRuntime,
  type KloAgentRuntimeDeps,
} from './agent-runtime.js';
import {
  isMissingProjectConfigError,
  missingConnectionSlSearchReadiness,
  missingProjectSlSearchReadiness,
  noConnectionsSlSearchReadiness,
  noIndexedSourcesSlSearchReadiness,
  type KloAgentSlSearchReadinessDetail,
} from './agent-search-readiness.js';
import { readKloSetupStatus, type KloSetupStatus } from './setup.js';

export type KloAgentArgs =
  | { command: 'tools'; projectDir: string; json: true }
  | { command: 'context'; projectDir: string; json: true }
  | { command: 'sl-list'; projectDir: string; json: true; connectionId?: string; query?: string }
  | { command: 'sl-read'; projectDir: string; json: true; connectionId?: string; sourceName: string }
  | {
      command: 'sl-query';
      projectDir: string;
      json: true;
      connectionId: string;
      queryFile: string;
      execute: boolean;
      maxRows?: number;
    }
  | { command: 'wiki-search'; projectDir: string; json: true; query: string; limit: number }
  | { command: 'wiki-read'; projectDir: string; json: true; pageId: string }
  | { command: 'sql-execute'; projectDir: string; json: true; connectionId: string; sqlFile: string; maxRows?: number };

export interface KloAgentDeps extends KloAgentRuntimeDeps {
  createRuntime?: (options: {
    projectDir: string;
    enableSemanticCompute: boolean;
    enableQueryExecution: boolean;
  }) => Promise<KloAgentRuntime>;
  readSetupStatus?: (
    projectDir: string,
  ) => Promise<KloSetupStatus | { project: { path?: string; ready: boolean }; agents: unknown[] }>;
}

const AGENT_TOOLS = [
  { name: 'context', command: 'klo agent context --json' },
  { name: 'sl.list', command: 'klo agent sl list --json [--connection-id <id>] [--query <text>]' },
  { name: 'sl.read', command: 'klo agent sl read <sourceName> --json [--connection-id <id>]' },
  {
    name: 'sl.query',
    command: 'klo agent sl query --json --connection-id <id> --query-file <path> --execute --max-rows 100',
  },
  { name: 'wiki.search', command: 'klo agent wiki search <query> --json [--limit 10]' },
  { name: 'wiki.read', command: 'klo agent wiki read <pageId> --json' },
  {
    name: 'sql.execute',
    command: 'klo agent sql execute --json --connection-id <id> --sql-file <path> --max-rows 100',
  },
] as const;

function writeAgentSlSearchReadinessError(io: KloCliIo, detail: KloAgentSlSearchReadinessDetail): void {
  writeAgentJsonError(io, detail.message, { code: detail.code, nextSteps: detail.nextSteps });
}

async function runtimeFor(args: KloAgentArgs, deps: KloAgentDeps): Promise<KloAgentRuntime> {
  const needsSemanticCompute = args.command === 'sl-query';
  const needsQueryExecution = args.command === 'sql-execute' || (args.command === 'sl-query' && args.execute);
  return deps.createRuntime
    ? deps.createRuntime({
        projectDir: args.projectDir,
        enableSemanticCompute: needsSemanticCompute,
        enableQueryExecution: needsQueryExecution,
      })
    : createKloAgentRuntime(
        {
          projectDir: args.projectDir,
          enableSemanticCompute: needsSemanticCompute,
          enableQueryExecution: needsQueryExecution,
        },
        deps,
      );
}

function connectionIdForSource(runtime: KloAgentRuntime, requested: string | undefined): string {
  if (requested) return requested;
  const ids = Object.keys(runtime.project.config.connections ?? {});
  if (ids.length === 1) return ids[0] as string;
  throw new Error('Use --connection-id when the project has zero or multiple connections.');
}

export async function runKloAgent(args: KloAgentArgs, io: KloCliIo, deps: KloAgentDeps = {}): Promise<number> {
  try {
    if (args.command === 'tools') {
      writeAgentJson(io, { projectDir: args.projectDir, tools: AGENT_TOOLS });
      return 0;
    }

    const runtime = await runtimeFor(args, deps);

    if (args.command === 'context') {
      const [status, connections, semanticLayer] = await Promise.all([
        (deps.readSetupStatus ?? readKloSetupStatus)(args.projectDir),
        runtime.ports.connections?.list() ?? [],
        runtime.ports.semanticLayer?.listSources({}) ?? { sources: [], totalSources: 0 },
      ]);
      writeAgentJson(io, { projectDir: args.projectDir, status, connections, semanticLayer, tools: AGENT_TOOLS });
      return 0;
    }

    if (args.command === 'sl-list') {
      const semanticLayer = runtime.ports.semanticLayer;
      if (!semanticLayer) throw new Error('Semantic-layer tools are not available for this project.');
      if (args.query) {
        const connectionIds = Object.keys(runtime.project.config.connections ?? {});
        if (args.connectionId && !runtime.project.config.connections[args.connectionId]) {
          writeAgentSlSearchReadinessError(
            io,
            missingConnectionSlSearchReadiness(args.projectDir, args.connectionId, args.query),
          );
          return 1;
        }
        if (connectionIds.length === 0) {
          writeAgentSlSearchReadinessError(io, noConnectionsSlSearchReadiness(args.projectDir, args.query));
          return 1;
        }
      }

      const listed = await semanticLayer.listSources({ connectionId: args.connectionId, query: args.query });
      if (args.query && listed.sources.length === 0) {
        const allSources = await semanticLayer.listSources({ connectionId: args.connectionId });
        if (allSources.totalSources === 0) {
          writeAgentSlSearchReadinessError(io, noIndexedSourcesSlSearchReadiness(args.projectDir, args.query));
          return 1;
        }
      }

      writeAgentJson(io, listed);
      return 0;
    }

    if (args.command === 'sl-read') {
      const semanticLayer = runtime.ports.semanticLayer;
      if (!semanticLayer) throw new Error('Semantic-layer tools are not available for this project.');
      const source = await semanticLayer.readSource({
        connectionId: connectionIdForSource(runtime, args.connectionId),
        sourceName: args.sourceName,
      });
      if (!source) throw new Error(`Semantic-layer source "${args.sourceName}" was not found.`);
      writeAgentJson(io, source);
      return 0;
    }

    if (args.command === 'sl-query') {
      const semanticLayer = runtime.ports.semanticLayer;
      if (!semanticLayer) throw new Error('Semantic-layer tools are not available for this project.');
      const query = await readAgentJsonFile(args.queryFile);
      const maxRows = args.execute ? parseAgentMaxRows(args.maxRows) : args.maxRows;
      writeAgentJson(
        io,
        await semanticLayer.query({
          connectionId: args.connectionId,
          query: { ...query, ...(maxRows !== undefined ? { limit: maxRows } : {}) } as never,
        }),
      );
      return 0;
    }

    if (args.command === 'wiki-search') {
      const knowledge = runtime.ports.knowledge;
      if (!knowledge) throw new Error('Wiki tools are not available for this project.');
      writeAgentJson(io, await knowledge.search({ userId: 'agent', query: args.query, limit: args.limit }));
      return 0;
    }

    if (args.command === 'wiki-read') {
      const knowledge = runtime.ports.knowledge;
      if (!knowledge) throw new Error('Wiki tools are not available for this project.');
      const page = await knowledge.read({ userId: 'agent', key: args.pageId });
      if (!page) throw new Error(`Wiki page "${args.pageId}" was not found.`);
      writeAgentJson(io, page);
      return 0;
    }

    const queryExecutor = runtime.queryExecutor;
    if (!queryExecutor) throw new Error('SQL execution is not available for this project.');
    const connection = runtime.project.config.connections[args.connectionId];
    if (!connection) throw new Error(`Connection "${args.connectionId}" was not found.`);
    const maxRows = parseAgentMaxRows(args.maxRows);
    writeAgentJson(
      io,
      await queryExecutor.execute({
        connectionId: args.connectionId,
        projectDir: runtime.project.projectDir,
        connection,
        sql: await readFile(args.sqlFile, 'utf-8'),
        maxRows,
      }),
    );
    return 0;
  } catch (error) {
    if (args.command === 'sl-list' && args.query && isMissingProjectConfigError(error)) {
      writeAgentSlSearchReadinessError(io, missingProjectSlSearchReadiness(args.projectDir, args.query));
      return 1;
    }
    writeAgentJsonError(io, error instanceof Error ? error.message : String(error));
    return 1;
  }
}