@agents-at-scale/ark 0.1.41 → 0.1.43

package/dist/lib/executeQuery.js

@@ -4,21 +4,91 @@
 import { execa } from 'execa';
 import ora from 'ora';
 import chalk from 'chalk';
-import output from './output.js';
 import { ExitCodes } from './errors.js';
-import { parseDuration } from './duration.js';
-/**
- * Execute a query against any ARK target (model, agent, team)
- * This is the shared implementation used by all query commands
- */
+import { ArkApiProxy } from './arkApiProxy.js';
+import { ChatClient } from './chatClient.js';
 export async function executeQuery(options) {
-    const spinner = ora('Creating query...').start();
-    const queryTimeoutMs = options.timeout
-        ? parseDuration(options.timeout)
-        : parseDuration('5m');
-    const watchTimeoutMs = options.watchTimeout
-        ? parseDuration(options.watchTimeout)
-        : queryTimeoutMs + 60000;
+    if (options.outputFormat) {
+        return executeQueryWithFormat(options);
+    }
+    let arkApiProxy;
+    const spinner = ora('Connecting to Ark API...').start();
+    try {
+        arkApiProxy = new ArkApiProxy();
+        const arkApiClient = await arkApiProxy.start();
+        const chatClient = new ChatClient(arkApiClient);
+        spinner.text = 'Executing query...';
+        const targetId = `${options.targetType}/${options.targetName}`;
+        const messages = [{ role: 'user', content: options.message }];
+        const state = {
+            toolCalls: new Map(),
+            content: '',
+        };
+        let lastAgentName;
+        let headerShown = false;
+        let firstOutput = true;
+        await chatClient.sendMessage(targetId, messages, { streamingEnabled: true }, (chunk, toolCalls, arkMetadata) => {
+            if (firstOutput) {
+                spinner.stop();
+                firstOutput = false;
+            }
+            const agentName = arkMetadata?.agent || arkMetadata?.team;
+            if (agentName && agentName !== lastAgentName) {
+                if (lastAgentName) {
+                    if (state.content) {
+                        process.stdout.write('\n');
+                    }
+                    process.stdout.write('\n');
+                }
+                const prefix = arkMetadata?.team ? '◆' : '●';
+                const color = arkMetadata?.team ? 'green' : 'cyan';
+                process.stdout.write(chalk[color](`${prefix} ${agentName}\n`));
+                lastAgentName = agentName;
+                state.content = '';
+                state.toolCalls.clear();
+                headerShown = true;
+            }
+            if (toolCalls && toolCalls.length > 0) {
+                for (const toolCall of toolCalls) {
+                    if (!state.toolCalls.has(toolCall.id)) {
+                        state.toolCalls.set(toolCall.id, toolCall);
+                        if (state.content) {
+                            process.stdout.write('\n');
+                        }
+                        process.stdout.write(chalk.magenta(`🔧 ${toolCall.function.name}\n`));
+                    }
+                }
+            }
+            if (chunk) {
+                if (state.toolCalls.size > 0 && !state.content) {
+                    process.stdout.write('\n');
+                }
+                process.stdout.write(chunk);
+                state.content += chunk;
+            }
+        });
+        if (spinner.isSpinning) {
+            spinner.stop();
+        }
+        if ((state.content || state.toolCalls.size > 0) && headerShown) {
+            process.stdout.write('\n');
+        }
+        if (arkApiProxy) {
+            arkApiProxy.stop();
+        }
+    }
+    catch (error) {
+        if (spinner.isSpinning) {
+            spinner.stop();
+        }
+        if (arkApiProxy) {
+            arkApiProxy.stop();
+        }
+        console.error(chalk.red(error instanceof Error ? error.message : 'Unknown error'));
+        process.exit(ExitCodes.CliError);
+    }
+}
+async function executeQueryWithFormat(options) {
     const timestamp = Date.now();
     const queryName = `cli-query-${timestamp}`;
     const queryManifest = {
@@ -39,71 +109,31 @@ export async function executeQuery(options) {
         },
     };
     try {
-        // Apply the query
-        spinner.text = 'Submitting query...';
         await execa('kubectl', ['apply', '-f', '-'], {
             input: JSON.stringify(queryManifest),
             stdio: ['pipe', 'pipe', 'pipe'],
         });
-        // Watch for query completion using kubectl wait
-        spinner.text = 'Waiting for query completion...';
-        try {
-            await execa('kubectl', [
-                'wait',
-                '--for=condition=Completed',
-                `query/${queryName}`,
-                `--timeout=${Math.floor(watchTimeoutMs / 1000)}s`,
-            ], { timeout: watchTimeoutMs });
+        const timeoutSeconds = 300;
+        await execa('kubectl', [
+            'wait',
+            '--for=condition=Completed',
+            `query/${queryName}`,
+            `--timeout=${timeoutSeconds}s`,
+        ], { timeout: timeoutSeconds * 1000 });
+        if (options.outputFormat === 'name') {
+            console.log(queryName);
         }
-        catch (error) {
-            spinner.stop();
-            // Check if it's a timeout or other error
-            if (error instanceof Error &&
-                error.message.includes('timed out waiting')) {
-                console.error(chalk.red(`Query did not complete within ${options.watchTimeout ?? `${Math.floor(watchTimeoutMs / 1000)}s`}`));
-                process.exit(ExitCodes.Timeout);
-            }
-            // For other errors, fetch the query to check status
-        }
-        spinner.stop();
-        // Fetch final query state
-        try {
-            const { stdout } = await execa('kubectl', ['get', 'query', queryName, '-o', 'json'], { stdio: 'pipe' });
-            const query = JSON.parse(stdout);
-            const phase = query.status?.phase;
-            // Check if query completed successfully or with error
-            if (phase === 'done') {
-                // Extract and display the response from responses array
-                if (query.status?.responses && query.status.responses.length > 0) {
-                    const response = query.status.responses[0];
-                    console.log(response.content || response);
-                }
-                else {
-                    output.warning('No response received');
-                }
-            }
-            else if (phase === 'error') {
-                const response = query.status?.responses?.[0];
-                console.error(chalk.red(response?.content || 'Query failed with unknown error'));
-                process.exit(ExitCodes.OperationError);
-            }
-            else if (phase === 'canceled') {
-                spinner.warn('Query canceled');
-                if (query.status?.message) {
-                    output.warning(query.status.message);
-                }
-                process.exit(ExitCodes.OperationError);
-            }
+        else if (options.outputFormat === 'json' ||
+            options.outputFormat === 'yaml') {
+            const { stdout } = await execa('kubectl', ['get', 'query', queryName, '-o', options.outputFormat], { stdio: 'pipe' });
+            console.log(stdout);
         }
-        catch (error) {
-            console.error(chalk.red(error instanceof Error
-                ? error.message
-                : 'Failed to fetch query result'));
+        else {
+            console.error(chalk.red(`Invalid output format: ${options.outputFormat}. Use: yaml, json, or name`));
             process.exit(ExitCodes.CliError);
         }
     }
     catch (error) {
-        spinner.stop();
         console.error(chalk.red(error instanceof Error ? error.message : 'Unknown error'));
         process.exit(ExitCodes.CliError);
     }

package/dist/lib/executeQuery.spec.js

@@ -10,17 +10,26 @@ const mockSpinner = {
     warn: jest.fn(),
     stop: jest.fn(),
     text: '',
+    isSpinning: false,
 };
 const mockOra = jest.fn(() => mockSpinner);
 jest.unstable_mockModule('ora', () => ({
     default: mockOra,
 }));
-const mockOutput = {
-    warning: jest.fn(),
-    error: jest.fn(),
+let mockSendMessage = jest.fn();
+const mockChatClient = jest.fn(() => ({
+    sendMessage: mockSendMessage,
+}));
+let mockArkApiProxyInstance = {
+    start: jest.fn(),
+    stop: jest.fn(),
 };
-jest.unstable_mockModule('./output.js', () => ({
-    default: mockOutput,
+const mockArkApiProxy = jest.fn(() => mockArkApiProxyInstance);
+jest.unstable_mockModule('./arkApiProxy.js', () => ({
+    ArkApiProxy: mockArkApiProxy,
+}));
+jest.unstable_mockModule('./chatClient.js', () => ({
+    ChatClient: mockChatClient,
 }));
 const mockExit = jest.spyOn(process, 'exit').mockImplementation((() => {
     throw new Error('process.exit called');
@@ -29,12 +38,25 @@ const mockConsoleLog = jest.spyOn(console, 'log').mockImplementation(() => { });
 const mockConsoleError = jest
     .spyOn(console, 'error')
     .mockImplementation(() => { });
+const mockStdoutWrite = jest
+    .spyOn(process.stdout, 'write')
+    .mockImplementation(() => true);
 const { executeQuery, parseTarget } = await import('./executeQuery.js');
 const { ExitCodes } = await import('./errors.js');
 describe('executeQuery', () => {
     beforeEach(() => {
         jest.clearAllMocks();
         mockSpinner.start.mockReturnValue(mockSpinner);
+        mockSpinner.isSpinning = false;
+        mockSendMessage = jest.fn();
+        mockChatClient.mockReturnValue({ sendMessage: mockSendMessage });
+        const startMock = jest.fn();
+        startMock.mockResolvedValue({});
+        mockArkApiProxyInstance = {
+            start: startMock,
+            stop: jest.fn(),
+        };
+        mockArkApiProxy.mockReturnValue(mockArkApiProxyInstance);
     });
     describe('parseTarget', () => {
         it('should parse valid target strings', () => {
@@ -57,24 +79,107 @@ describe('executeQuery', () => {
             expect(parseTarget('model/default/extra')).toBeNull();
         });
     });
-    describe('executeQuery', () => {
-        it('should create and apply a query manifest', async () => {
-            const mockQueryResponse = {
-                status: {
-                    phase: 'done',
-                    responses: [{ content: 'Test response' }],
-                },
-            };
+    describe('executeQuery with streaming', () => {
+        it('should execute query with streaming and display chunks', async () => {
+            mockSendMessage.mockImplementation(async (targetId, messages, options, callback) => {
+                callback('Hello', undefined, { agent: 'test-agent' });
+                callback(' world', undefined, { agent: 'test-agent' });
+            });
+            await executeQuery({
+                targetType: 'model',
+                targetName: 'default',
+                message: 'Hello',
+            });
+            expect(mockArkApiProxy).toHaveBeenCalled();
+            expect(mockArkApiProxyInstance.start).toHaveBeenCalled();
+            expect(mockChatClient).toHaveBeenCalled();
+            expect(mockSendMessage).toHaveBeenCalledWith('model/default', [{ role: 'user', content: 'Hello' }], { streamingEnabled: true }, expect.any(Function));
+            expect(mockSpinner.stop).toHaveBeenCalled();
+            expect(mockArkApiProxyInstance.stop).toHaveBeenCalled();
+            expect(mockStdoutWrite).toHaveBeenCalled();
+        });
+        it('should display agent names with correct formatting', async () => {
+            mockSendMessage.mockImplementation(async (targetId, messages, options, callback) => {
+                callback('Response 1', undefined, { agent: 'agent-1' });
+                callback('Response 2', undefined, { agent: 'agent-2' });
+            });
+            await executeQuery({
+                targetType: 'agent',
+                targetName: 'test-agent',
+                message: 'Hello',
+            });
+            expect(mockStdoutWrite).toHaveBeenCalled();
+            const calls = mockStdoutWrite.mock.calls.map((call) => String(call[0]));
+            expect(calls.some((call) => call.includes('agent-1'))).toBe(true);
+            expect(calls.some((call) => call.includes('agent-2'))).toBe(true);
+        });
+        it('should display team names with diamond prefix', async () => {
+            mockSendMessage.mockImplementation(async (targetId, messages, options, callback) => {
+                callback('Team response', undefined, { team: 'my-team' });
+            });
+            await executeQuery({
+                targetType: 'team',
+                targetName: 'my-team',
+                message: 'Hello',
+            });
+            const calls = mockStdoutWrite.mock.calls.map((call) => String(call[0]));
+            expect(calls.some((call) => call.includes('◆'))).toBe(true);
+            expect(calls.some((call) => call.includes('my-team'))).toBe(true);
+        });
+        it('should display tool calls', async () => {
+            mockSendMessage.mockImplementation(async (targetId, messages, options, callback) => {
+                callback('', [{ id: 1, function: { name: 'get_weather' } }], {
+                    agent: 'weather-agent',
+                });
+                callback('The weather is sunny', undefined, {
+                    agent: 'weather-agent',
+                });
+            });
+            await executeQuery({
+                targetType: 'agent',
+                targetName: 'weather-agent',
+                message: 'What is the weather?',
+            });
+            const calls = mockStdoutWrite.mock.calls.map((call) => String(call[0]));
+            expect(calls.some((call) => call.includes('get_weather'))).toBe(true);
+            expect(calls.some((call) => call.includes('The weather is sunny'))).toBe(true);
+        });
+        it('should handle errors and exit with CliError', async () => {
+            mockSpinner.isSpinning = true;
+            const startMock = jest.fn();
+            startMock.mockRejectedValue(new Error('Connection failed'));
+            mockArkApiProxyInstance.start = startMock;
+            await expect(executeQuery({
+                targetType: 'model',
+                targetName: 'default',
+                message: 'Hello',
+            })).rejects.toThrow('process.exit called');
+            expect(mockSpinner.stop).toHaveBeenCalled();
+            expect(mockConsoleError).toHaveBeenCalledWith(expect.stringContaining('Connection failed'));
+            expect(mockExit).toHaveBeenCalledWith(ExitCodes.CliError);
+            expect(mockArkApiProxyInstance.stop).toHaveBeenCalled();
+        });
+        it('should stop spinner when first output arrives', async () => {
+            mockSpinner.isSpinning = true;
+            mockSendMessage.mockImplementation(async (targetId, messages, options, callback) => {
+                callback('First chunk', undefined, { agent: 'test-agent' });
+            });
+            await executeQuery({
+                targetType: 'model',
+                targetName: 'default',
+                message: 'Hello',
+            });
+            expect(mockSpinner.stop).toHaveBeenCalled();
+        });
+    });
+    describe('executeQuery with output format', () => {
+        it('should create query and output name format', async () => {
             mockExeca.mockImplementation(async (command, args) => {
                 if (args.includes('apply')) {
                     return { stdout: '', stderr: '', exitCode: 0 };
                 }
-                if (args.includes('get') && args.includes('query')) {
-                    return {
-                        stdout: JSON.stringify(mockQueryResponse),
-                        stderr: '',
-                        exitCode: 0,
-                    };
+                if (args.includes('wait')) {
+                    return { stdout: '', stderr: '', exitCode: 0 };
                 }
                 return { stdout: '', stderr: '', exitCode: 0 };
             });
@@ -82,83 +187,67 @@ describe('executeQuery', () => {
                 targetType: 'model',
                 targetName: 'default',
                 message: 'Hello',
+                outputFormat: 'name',
             });
-            expect(mockSpinner.start).toHaveBeenCalled();
-            expect(mockSpinner.stop).toHaveBeenCalled();
-            expect(mockConsoleLog).toHaveBeenCalledWith('Test response');
+            expect(mockExeca).toHaveBeenCalledWith('kubectl', expect.arrayContaining(['apply', '-f', '-']), expect.any(Object));
+            expect(mockExeca).toHaveBeenCalledWith('kubectl', expect.arrayContaining(['wait', '--for=condition=Completed']), expect.any(Object));
+            expect(mockConsoleLog).toHaveBeenCalledWith(expect.stringMatching(/cli-query-\d+/));
         });
-        it('should handle query error phase and exit with code 2', async () => {
-            const mockQueryResponse = {
-                status: {
-                    phase: 'error',
-                    responses: [{ content: 'Query failed with test error' }],
-                },
+        it('should output json format', async () => {
+            const mockQuery = {
+                apiVersion: 'ark.mckinsey.com/v1alpha1',
+                kind: 'Query',
+                metadata: { name: 'test-query' },
             };
             mockExeca.mockImplementation(async (command, args) => {
                 if (args.includes('apply')) {
                     return { stdout: '', stderr: '', exitCode: 0 };
                 }
-                if (args.includes('get') && args.includes('query')) {
-                    return {
-                        stdout: JSON.stringify(mockQueryResponse),
-                        stderr: '',
-                        exitCode: 0,
-                    };
+                if (args.includes('wait')) {
+                    return { stdout: '', stderr: '', exitCode: 0 };
+                }
+                if (args.includes('get') && args.includes('-o')) {
+                    return { stdout: JSON.stringify(mockQuery), stderr: '', exitCode: 0 };
                 }
                 return { stdout: '', stderr: '', exitCode: 0 };
             });
-            try {
-                await executeQuery({
-                    targetType: 'model',
-                    targetName: 'default',
-                    message: 'Hello',
-                });
-            }
-            catch (error) {
-                expect(error.message).toBe('process.exit called');
-            }
-            expect(mockSpinner.stop).toHaveBeenCalled();
-            expect(mockConsoleError).toHaveBeenCalledWith(expect.stringContaining('Query failed with test error'));
-            expect(mockExit).toHaveBeenCalledWith(ExitCodes.OperationError);
+            await executeQuery({
+                targetType: 'model',
+                targetName: 'default',
+                message: 'Hello',
+                outputFormat: 'json',
+            });
+            expect(mockConsoleLog).toHaveBeenCalledWith(JSON.stringify(mockQuery));
         });
-        it('should handle query canceled phase and exit with code 2', async () => {
-            const mockQueryResponse = {
-                status: {
-                    phase: 'canceled',
-                    message: 'Query was canceled',
-                },
-            };
+        it('should output yaml format', async () => {
+            const mockYaml = 'apiVersion: ark.mckinsey.com/v1alpha1\nkind: Query';
             mockExeca.mockImplementation(async (command, args) => {
                 if (args.includes('apply')) {
                     return { stdout: '', stderr: '', exitCode: 0 };
                 }
-                if (args.includes('get') && args.includes('query')) {
-                    return {
-                        stdout: JSON.stringify(mockQueryResponse),
-                        stderr: '',
-                        exitCode: 0,
-                    };
+                if (args.includes('wait')) {
+                    return { stdout: '', stderr: '', exitCode: 0 };
+                }
+                if (args.includes('get') && args.includes('yaml')) {
+                    return { stdout: mockYaml, stderr: '', exitCode: 0 };
                 }
                 return { stdout: '', stderr: '', exitCode: 0 };
             });
-            try {
-                await executeQuery({
-                    targetType: 'agent',
-                    targetName: 'test-agent',
-                    message: 'Hello',
-                });
-            }
-            catch (error) {
-                expect(error.message).toBe('process.exit called');
-            }
-            expect(mockSpinner.warn).toHaveBeenCalledWith('Query canceled');
-            expect(mockOutput.warning).toHaveBeenCalledWith('Query was canceled');
-            expect(mockExit).toHaveBeenCalledWith(ExitCodes.OperationError);
+            await executeQuery({
+                targetType: 'model',
+                targetName: 'default',
+                message: 'Hello',
+                outputFormat: 'yaml',
+            });
+            expect(mockConsoleLog).toHaveBeenCalledWith(mockYaml);
         });
-        it('should handle kubectl apply failures with exit code 1', async () => {
+        it('should reject invalid output format', async () => {
             mockExeca.mockImplementation(async (command, args) => {
                 if (args.includes('apply')) {
-                    throw new Error('Failed to apply');
+                    return { stdout: '', stderr: '', exitCode: 0 };
+                }
+                if (args.includes('wait')) {
+                    return { stdout: '', stderr: '', exitCode: 0 };
                 }
                 return { stdout: '', stderr: '', exitCode: 0 };
             });
@@ -166,38 +255,26 @@ describe('executeQuery', () => {
                 targetType: 'model',
                 targetName: 'default',
                 message: 'Hello',
+                outputFormat: 'invalid',
             })).rejects.toThrow('process.exit called');
-            expect(mockSpinner.stop).toHaveBeenCalled();
-            expect(mockConsoleError).toHaveBeenCalledWith(expect.stringContaining('Failed to apply'));
+            expect(mockConsoleError).toHaveBeenCalledWith(expect.stringContaining('Invalid output format'));
             expect(mockExit).toHaveBeenCalledWith(ExitCodes.CliError);
         });
-        it('should handle query timeout and exit with code 3', async () => {
+        it('should handle kubectl errors', async () => {
             mockExeca.mockImplementation(async (command, args) => {
                 if (args.includes('apply')) {
-                    return { stdout: '', stderr: '', exitCode: 0 };
-                }
-                if (args.includes('wait')) {
-                    // Simulate kubectl wait timeout
-                    const error = new Error('timed out waiting for the condition');
-                    throw error;
+                    throw new Error('kubectl apply failed');
                 }
                 return { stdout: '', stderr: '', exitCode: 0 };
             });
-            try {
-                await executeQuery({
-                    targetType: 'model',
-                    targetName: 'default',
-                    message: 'Hello',
-                    timeout: '100ms',
-                    watchTimeout: '200ms',
-                });
-            }
-            catch (error) {
-                expect(error.message).toBe('process.exit called');
-            }
-            expect(mockSpinner.stop).toHaveBeenCalled();
-            expect(mockConsoleError).toHaveBeenCalledWith(expect.stringContaining('Query did not complete within 200ms'));
-            expect(mockExit).toHaveBeenCalledWith(ExitCodes.Timeout);
+            await expect(executeQuery({
+                targetType: 'model',
+                targetName: 'default',
+                message: 'Hello',
+                outputFormat: 'name',
+            })).rejects.toThrow('process.exit called');
+            expect(mockConsoleError).toHaveBeenCalledWith(expect.stringContaining('kubectl apply failed'));
+            expect(mockExit).toHaveBeenCalledWith(ExitCodes.CliError);
         });
     });
 });

package/dist/lib/kubectl.d.ts

@@ -0,0 +1,15 @@
+interface K8sResource {
+    metadata: {
+        name: string;
+        creationTimestamp?: string;
+    };
+}
+export declare function getResource<T extends K8sResource>(resourceType: string, name: string): Promise<T>;
+export declare function listResources<T extends K8sResource>(resourceType: string, options?: {
+    sortBy?: string;
+}): Promise<T[]>;
+export declare function deleteResource(resourceType: string, name?: string, options?: {
+    all?: boolean;
+}): Promise<void>;
+export declare function replaceResource<T extends K8sResource>(resource: T): Promise<T>;
+export {};

package/dist/lib/kubectl.js

@@ -0,0 +1,47 @@
+import { execa } from 'execa';
+export async function getResource(resourceType, name) {
+    if (name === '@latest') {
+        const result = await execa('kubectl', [
+            'get',
+            resourceType,
+            '--sort-by=.metadata.creationTimestamp',
+            '-o',
+            'json',
+        ], { stdio: 'pipe' });
+        const data = JSON.parse(result.stdout);
+        const resources = data.items || [];
+        if (resources.length === 0) {
+            throw new Error(`No ${resourceType} found`);
+        }
+        return resources[resources.length - 1];
+    }
+    const result = await execa('kubectl', ['get', resourceType, name, '-o', 'json'], { stdio: 'pipe' });
+    return JSON.parse(result.stdout);
+}
+export async function listResources(resourceType, options) {
+    const args = ['get', resourceType];
+    if (options?.sortBy) {
+        args.push(`--sort-by=${options.sortBy}`);
+    }
+    args.push('-o', 'json');
+    const result = await execa('kubectl', args, { stdio: 'pipe' });
+    const data = JSON.parse(result.stdout);
+    return data.items || [];
+}
+export async function deleteResource(resourceType, name, options) {
+    const args = ['delete', resourceType];
+    if (options?.all) {
+        args.push('--all');
+    }
+    else if (name) {
+        args.push(name);
+    }
+    await execa('kubectl', args, { stdio: 'pipe' });
+}
+export async function replaceResource(resource) {
+    const result = await execa('kubectl', ['replace', '-f', '-', '-o', 'json'], {
+        input: JSON.stringify(resource),
+        stdio: ['pipe', 'pipe', 'pipe'],
+    });
+    return JSON.parse(result.stdout);
+}

package/dist/lib/kubectl.spec.d.ts

@@ -0,0 +1 @@
+export {};