diff --git a/package.json b/package.json index a5e9ae9..bda03ba 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@kitops/kitops-ts", - "version": "0.0.4", + "version": "0.1.0", "description": "TypeScript library for KitOps CLI", "type": "module", "main": "./dist/index.js", diff --git a/src/commands/__tests__/pack.spec.ts b/src/commands/__tests__/pack.spec.ts index 22d3a36..8641305 100644 --- a/src/commands/__tests__/pack.spec.ts +++ b/src/commands/__tests__/pack.spec.ts @@ -2,7 +2,10 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { pack } from '../pack'; import { runCommand, prepareArgs } from '../../core/exec'; -vi.mock('../../core/exec'); +vi.mock('../../core/exec', async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, runCommand: vi.fn(), prepareArgs: vi.fn() }; +}); const mockRunCommand = vi.mocked(runCommand); const mockPrepareArgs = vi.mocked(prepareArgs); @@ -10,7 +13,6 @@ const mockPrepareArgs = vi.mocked(prepareArgs); describe('pack', () => { beforeEach(() => { vi.clearAllMocks(); - // Mock prepareArgs to return the expected array format mockPrepareArgs.mockImplementation((options) => { const args: string[] = []; Object.entries(options).forEach(([key, value]) => { @@ -22,39 +24,29 @@ describe('pack', () => { }); return args; }); + mockRunCommand.mockResolvedValue({ stdout: 'Pack completed successfully', stderr: '', exitCode: 0 }); }); - it('should call runCommand with default directory when no arguments provided', async () => { - mockRunCommand.mockResolvedValue({ - stdout: 'Pack completed successfully', - stderr: '', - exitCode: 0 - }); + it('should return a CancellablePromise with a cancel method', () => { + const op = pack(); + expect(op).toBeInstanceOf(Promise); + expect(typeof op.cancel).toBe('function'); + return op; + }); + it('should call runCommand with default directory when no arguments provided', async () => { await pack(); - expect(mockRunCommand).toHaveBeenCalledWith('pack', ['.']); + expect(mockRunCommand).toHaveBeenCalledWith('pack', ['.'], undefined, { signal: expect.any(AbortSignal) }); }); it('should call runCommand with specified directory', async () => { - mockRunCommand.mockResolvedValue({ - stdout: 'Pack completed successfully', - stderr: '', - exitCode: 0 - }); - await pack('/path/to/project'); - expect(mockRunCommand).toHaveBeenCalledWith('pack', ['/path/to/project']); + expect(mockRunCommand).toHaveBeenCalledWith('pack', ['/path/to/project'], undefined, { signal: expect.any(AbortSignal) }); }); it('should handle pack with all flags', async () => { - mockRunCommand.mockResolvedValue({ - stdout: 'Pack completed successfully', - stderr: '', - exitCode: 0 - }); - const flags = { file: 'Kitfile', tag: 'my-model:v1.0.0', @@ -71,16 +63,10 @@ describe('pack', () => { '--tag', 'my-model:v1.0.0', '--compression', 'gzip', '--useModelPack' - ]); + ], undefined, { signal: expect.any(AbortSignal) }); }); it('should handle pack with partial flags', async () => { - mockRunCommand.mockResolvedValue({ - stdout: 'Pack completed successfully', - stderr: '', - exitCode: 0 - }); - const flags = { file: 'CustomKitfile', tag: 'test-model:latest', @@ -96,16 +82,10 @@ describe('pack', () => { '--file', 'CustomKitfile', '--tag', 'test-model:latest', '--compression', 'none' - ]); + ], undefined, { signal: expect.any(AbortSignal) }); }); it('should handle pack with only file flag', async () => { - mockRunCommand.mockResolvedValue({ - stdout: 'Pack completed successfully', - stderr: '', - exitCode: 0 - }); - const flags = { file: 'MyKitfile', tag: '', @@ -113,7 +93,6 @@ describe('pack', () => { useModelPack: false }; - // Mock prepareArgs to handle empty strings appropriately mockPrepareArgs.mockReturnValue(['--file', 'MyKitfile']); await pack('./models', flags); @@ -122,16 +101,10 @@ describe('pack', () => { expect(mockRunCommand).toHaveBeenCalledWith('pack', [ './models', '--file', 'MyKitfile' - ]); + ], undefined, { signal: expect.any(AbortSignal) }); }); it('should handle pack with only useModelPack flag', async () => { - mockRunCommand.mockResolvedValue({ - stdout: 'Pack completed successfully', - stderr: '', - exitCode: 0 - }); - const flags = { file: '', tag: '', @@ -147,33 +120,21 @@ describe('pack', () => { expect(mockRunCommand).toHaveBeenCalledWith('pack', [ '.', '--useModelPack' - ]); + ], undefined, { signal: expect.any(AbortSignal) }); }); it('should propagate errors from runCommand', async () => { const errorMessage = 'Kit command failed with exit code 1: Kitfile not found'; mockRunCommand.mockRejectedValue(new Error(errorMessage)); - await expect(pack('./nonexistent')) - .rejects.toThrow(errorMessage); + await expect(pack('./nonexistent')).rejects.toThrow(errorMessage); }); it('should handle different compression types', async () => { - mockRunCommand.mockResolvedValue({ - stdout: 'Pack completed successfully', - stderr: '', - exitCode: 0 - }); - const compressionTypes = ['gzip', 'zstd', 'none']; for (const compression of compressionTypes) { - const flags = { - file: 'Kitfile', - tag: 'model:latest', - compression, - useModelPack: false - }; + const flags = { file: 'Kitfile', tag: 'model:latest', compression, useModelPack: false }; mockPrepareArgs.mockReturnValue([ '--file', 'Kitfile', @@ -188,17 +149,11 @@ describe('pack', () => { '--file', 'Kitfile', '--tag', 'model:latest', '--compression', compression - ]); + ], undefined, { signal: expect.any(AbortSignal) }); } }); it('should handle different tag formats', async () => { - mockRunCommand.mockResolvedValue({ - stdout: 'Pack completed successfully', - stderr: '', - exitCode: 0 - }); - const tagFormats = [ 'simple-model:v1.0.0', 'registry.example.com/org/model:latest', @@ -208,12 +163,7 @@ describe('pack', () => { ]; for (const tag of tagFormats) { - const flags = { - file: 'Kitfile', - tag, - compression: 'gzip', - useModelPack: false - }; + const flags = { file: 'Kitfile', tag, compression: 'gzip', useModelPack: false }; mockPrepareArgs.mockReturnValue([ '--file', 'Kitfile', @@ -228,34 +178,21 @@ describe('pack', () => { '--file', 'Kitfile', '--tag', tag, '--compression', 'gzip' - ]); + ], undefined, { signal: expect.any(AbortSignal) }); } }); it('should handle relative and absolute directory paths', async () => { - mockRunCommand.mockResolvedValue({ - stdout: 'Pack completed successfully', - stderr: '', - exitCode: 0 - }); - - const directories = [ - '.', - './src', - '../project', - '/absolute/path/to/project', - '~/home/user/project' - ]; + const directories = ['.', './src', '../project', '/absolute/path/to/project', '~/home/user/project']; for (const directory of directories) { await pack(directory); - expect(mockRunCommand).toHaveBeenCalledWith('pack', [directory]); + expect(mockRunCommand).toHaveBeenCalledWith('pack', [directory], undefined, { signal: expect.any(AbortSignal) }); } }); it('should handle pack command execution failure', async () => { - const execError = 'Failed to execute kit command: permission denied'; - mockRunCommand.mockRejectedValue(new Error(execError)); + mockRunCommand.mockRejectedValue(new Error('Failed to execute kit command: permission denied')); await expect(pack('./restricted-dir')) .rejects.toThrow('Failed to execute kit command: permission denied'); @@ -265,31 +202,28 @@ describe('pack', () => { const kitfileError = 'Kit command failed with exit code 1: invalid Kitfile path'; mockRunCommand.mockRejectedValue(new Error(kitfileError)); - const flags = { - file: '/invalid/path/Kitfile', - tag: 'model:latest', - compression: 'gzip', - useModelPack: false - }; + const flags = { file: '/invalid/path/Kitfile', tag: 'model:latest', compression: 'gzip', useModelPack: false }; - await expect(pack('.', flags)) - .rejects.toThrow('Kit command failed with exit code 1: invalid Kitfile path'); + await expect(pack('.', flags)).rejects.toThrow('Kit command failed with exit code 1: invalid Kitfile path'); }); it('should handle successful pack operation', async () => { - mockRunCommand.mockResolvedValue({ - stdout: 'Successfully packed model:v1.0.0', - stderr: '', - exitCode: 0 - }); - - const flags = { - file: 'Kitfile', - tag: 'model:v1.0.0', - compression: 'gzip', - useModelPack: true - }; + const flags = { file: 'Kitfile', tag: 'model:v1.0.0', compression: 'gzip', useModelPack: true }; await expect(pack('./project', flags)).resolves.not.toThrow(); }); -}); \ No newline at end of file + + it('should abort the signal when cancel is called', () => { + let capturedSignal: AbortSignal | undefined; + + mockRunCommand.mockImplementation((_cmd, _args, _stdin, opts) => { + capturedSignal = opts?.signal; + return Promise.resolve({ stdout: '', stderr: '', exitCode: 0 }); + }); + + const op = pack('.'); + expect(capturedSignal?.aborted).toBe(false); + op.cancel(); + expect(capturedSignal?.aborted).toBe(true); + }); +}); diff --git a/src/commands/__tests__/pull.spec.ts b/src/commands/__tests__/pull.spec.ts index 9358675..fe971cf 100644 --- a/src/commands/__tests__/pull.spec.ts +++ b/src/commands/__tests__/pull.spec.ts @@ -2,7 +2,10 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { pull } from '../pull'; import { runCommand, prepareArgs } from '../../core/exec'; -vi.mock('../../core/exec'); +vi.mock('../../core/exec', async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, runCommand: vi.fn(), prepareArgs: vi.fn() }; +}); const mockRunCommand = vi.mocked(runCommand); const mockPrepareArgs = vi.mocked(prepareArgs); @@ -14,10 +17,22 @@ describe('pull', () => { mockRunCommand.mockResolvedValue({ stdout: '', stderr: '', exitCode: 0 }); }); - it('should call runCommand with the reference', async () => { + it('should return a CancellablePromise with a cancel method', () => { + const op = pull('registry.example.com/org/my-model:v1.0.0'); + expect(op).toBeInstanceOf(Promise); + expect(typeof op.cancel).toBe('function'); + return op; + }); + + it('should call runCommand with the reference and an AbortSignal', async () => { await pull('registry.example.com/org/my-model:v1.0.0'); - expect(mockRunCommand).toHaveBeenCalledWith('pull', ['registry.example.com/org/my-model:v1.0.0']); + expect(mockRunCommand).toHaveBeenCalledWith( + 'pull', + ['registry.example.com/org/my-model:v1.0.0'], + undefined, + { signal: expect.any(AbortSignal) }, + ); }); it('should forward TLS flags', async () => { @@ -25,16 +40,23 @@ describe('pull', () => { await pull('registry.example.com/org/my-model:v1.0.0', { tlsCert: '/path/to/cert.pem' }); - expect(mockRunCommand).toHaveBeenCalledWith('pull', [ - 'registry.example.com/org/my-model:v1.0.0', - '--tls-cert=/path/to/cert.pem', - ]); + expect(mockRunCommand).toHaveBeenCalledWith( + 'pull', + ['registry.example.com/org/my-model:v1.0.0', '--tls-cert=/path/to/cert.pem'], + undefined, + { signal: expect.any(AbortSignal) }, + ); }); it('should work without flags', async () => { await pull('my-model:latest'); - expect(mockRunCommand).toHaveBeenCalledWith('pull', ['my-model:latest']); + expect(mockRunCommand).toHaveBeenCalledWith( + 'pull', + ['my-model:latest'], + undefined, + { signal: expect.any(AbortSignal) }, + ); }); it('should propagate errors from runCommand', async () => { @@ -43,4 +65,18 @@ describe('pull', () => { await expect(pull('registry.example.com/org/my-model:v1.0.0')) .rejects.toThrow('Kit command failed with exit code 1: not found'); }); + + it('should abort the signal when cancel is called', () => { + let capturedSignal: AbortSignal | undefined; + + mockRunCommand.mockImplementation((_cmd, _args, _stdin, opts) => { + capturedSignal = opts?.signal; + return Promise.resolve({ stdout: '', stderr: '', exitCode: 0 }); + }); + + const op = pull('registry.example.com/org/my-model:v1.0.0'); + expect(capturedSignal?.aborted).toBe(false); + op.cancel(); + expect(capturedSignal?.aborted).toBe(true); + }); }); diff --git a/src/commands/__tests__/unpack.spec.ts b/src/commands/__tests__/unpack.spec.ts index e7389cf..e366907 100644 --- a/src/commands/__tests__/unpack.spec.ts +++ b/src/commands/__tests__/unpack.spec.ts @@ -2,7 +2,10 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { unpack } from '../unpack'; import { runCommand, prepareArgs } from '../../core/exec'; -vi.mock('../../core/exec'); +vi.mock('../../core/exec', async (importOriginal) => { + const actual = await importOriginal(); + return { ...actual, runCommand: vi.fn(), prepareArgs: vi.fn() }; +}); const mockRunCommand = vi.mocked(runCommand); const mockPrepareArgs = vi.mocked(prepareArgs); @@ -14,10 +17,17 @@ describe('unpack', () => { mockRunCommand.mockResolvedValue({ stdout: '', stderr: '', exitCode: 0 }); }); - it('should call runCommand with the destination path', async () => { + it('should return a CancellablePromise with a cancel method', () => { + const op = unpack('./output'); + expect(op).toBeInstanceOf(Promise); + expect(typeof op.cancel).toBe('function'); + return op; + }); + + it('should call runCommand with the destination path and an AbortSignal', async () => { await unpack('./output'); - expect(mockRunCommand).toHaveBeenCalledWith('unpack', ['./output']); + expect(mockRunCommand).toHaveBeenCalledWith('unpack', ['./output'], undefined, { signal: expect.any(AbortSignal) }); }); it('should forward flags to prepareArgs', async () => { @@ -26,7 +36,7 @@ describe('unpack', () => { await unpack('./output', { filter: 'model' }); expect(mockPrepareArgs).toHaveBeenCalledWith({ filter: 'model' }); - expect(mockRunCommand).toHaveBeenCalledWith('unpack', ['./output', '--filter=model']); + expect(mockRunCommand).toHaveBeenCalledWith('unpack', ['./output', '--filter=model'], undefined, { signal: expect.any(AbortSignal) }); }); it('should support overwrite flag', async () => { @@ -34,7 +44,7 @@ describe('unpack', () => { await unpack('./output', { overwrite: true }); - expect(mockRunCommand).toHaveBeenCalledWith('unpack', ['./output', '--overwrite']); + expect(mockRunCommand).toHaveBeenCalledWith('unpack', ['./output', '--overwrite'], undefined, { signal: expect.any(AbortSignal) }); }); it('should support ignoreExisting flag', async () => { @@ -42,7 +52,7 @@ describe('unpack', () => { await unpack('./output', { ignoreExisting: true }); - expect(mockRunCommand).toHaveBeenCalledWith('unpack', ['./output', '--ignore-existing']); + expect(mockRunCommand).toHaveBeenCalledWith('unpack', ['./output', '--ignore-existing'], undefined, { signal: expect.any(AbortSignal) }); }); it('should propagate errors from runCommand', async () => { @@ -51,4 +61,18 @@ describe('unpack', () => { await expect(unpack('./output')) .rejects.toThrow('Kit command failed with exit code 1: permission denied'); }); + + it('should abort the signal when cancel is called', () => { + let capturedSignal: AbortSignal | undefined; + + mockRunCommand.mockImplementation((_cmd, _args, _stdin, opts) => { + capturedSignal = opts?.signal; + return Promise.resolve({ stdout: '', stderr: '', exitCode: 0 }); + }); + + const op = unpack('./output'); + expect(capturedSignal?.aborted).toBe(false); + op.cancel(); + expect(capturedSignal?.aborted).toBe(true); + }); }); diff --git a/src/commands/pack.ts b/src/commands/pack.ts index 0dc8606..4cedf35 100644 --- a/src/commands/pack.ts +++ b/src/commands/pack.ts @@ -1,5 +1,6 @@ -import { runCommand, prepareArgs } from "../core/exec.js"; +import { runCommand, prepareArgs, cancellable } from "../core/exec.js"; import type { PackFlags } from "../types/commands.js"; +import type { CancellablePromise } from "../types/kitops.js"; /** * Packages a ModelKit from a directory that contains a Kitfile. @@ -7,17 +8,30 @@ import type { PackFlags } from "../types/commands.js"; * Use `flags.tag` to name the resulting ModelKit so it can be pushed directly. * If no tag is provided the kit is stored locally without a named reference. * + * Returns a {@link CancellablePromise}. Call `.cancel()` on the returned value to abort + * the operation and kill the underlying `kit` process at any time. + * * @param directory - Directory containing the Kitfile (defaults to the current directory). * @see https://kitops.org/docs/cli/cli-reference/#kit-pack + * + * @example + * ```ts + * const op = pack('.', { tag: 'registry.example.com/org/model:v1' }); + * setTimeout(() => op.cancel(), 60_000); + * try { + * await op; + * } catch (e) { + * if (e instanceof DOMException && e.name === 'AbortError') { + * console.log('Pack was cancelled'); + * } + * } + * ``` */ -export async function pack(directory: string = '.', flags?: PackFlags): Promise { - const args = [directory] - - if (flags) { - args.push(...prepareArgs(flags)) - } - - await runCommand('pack', args); +export function pack(directory: string = '.', flags?: PackFlags): CancellablePromise { + return cancellable((signal) => { + const args = [directory, ...(flags ? prepareArgs(flags) : [])]; + return runCommand('pack', args, undefined, { signal }).then(() => {}); + }); // @TODO: return pack result (tag, digest) and any other useful info -} \ No newline at end of file +} diff --git a/src/commands/pull.ts b/src/commands/pull.ts index 72b0f46..4c6e6cd 100644 --- a/src/commands/pull.ts +++ b/src/commands/pull.ts @@ -1,5 +1,6 @@ -import { prepareArgs, runCommand } from "../core/exec.js"; -import { TLSFlags } from "../types/commands.js"; +import { prepareArgs, runCommand, cancellable } from "../core/exec.js"; +import type { TLSFlags } from "../types/commands.js"; +import type { CancellablePromise } from "../types/kitops.js"; /** * Pulls a ModelKit from a registry into local storage. @@ -7,16 +8,29 @@ import { TLSFlags } from "../types/commands.js"; * After pulling, use `unpack` to extract the contents to disk or * `inspect` / `info` to read its metadata without extracting. * + * Returns a {@link CancellablePromise}. Call `.cancel()` on the returned value to abort + * the transfer and kill the underlying `kit` process at any time. + * * @param path - Full ModelKit reference path in the form of `registry/repository[:tag|@digest]`. * @param flags - Optional flags to modify the pull behavior (e.g. TLS settings). * @see https://kitops.org/docs/cli/cli-reference/#kit-pull + * + * @example + * ```ts + * const op = pull('registry.example.com/org/model:v1'); + * setTimeout(() => op.cancel(), 30_000); + * try { + * await op; + * } catch (e) { + * if (e instanceof DOMException && e.name === 'AbortError') { + * console.log('Pull was cancelled'); + * } + * } + * ``` */ -export async function pull(path: string, flags?: TLSFlags): Promise { - const args = [path]; - - if (flags) { - args.push(...prepareArgs(flags)) - } - - await runCommand('pull', args); -} \ No newline at end of file +export function pull(path: string, flags?: TLSFlags): CancellablePromise { + return cancellable((signal) => { + const args = [path, ...(flags ? prepareArgs(flags) : [])]; + return runCommand('pull', args, undefined, { signal }).then(() => {}); + }); +} diff --git a/src/commands/unpack.ts b/src/commands/unpack.ts index ba1843d..4b8fed6 100644 --- a/src/commands/unpack.ts +++ b/src/commands/unpack.ts @@ -1,18 +1,32 @@ -import { runCommand, prepareArgs } from "../core/exec.js"; +import { runCommand, prepareArgs, cancellable } from "../core/exec.js"; import type { UnpackFlags } from "../types/commands.js"; +import type { CancellablePromise } from "../types/kitops.js"; /** - * Extracts a ModelKit with the given flags + * Extracts a ModelKit with the given flags. + * + * Returns a {@link CancellablePromise}. Call `.cancel()` on the returned value to abort + * the extraction and kill the underlying `kit` process at any time. * * @param path - Full ModelKit reference path in the form of `registry/repository[:tag|@digest]`. * @see https://kitops.org/docs/cli/cli-reference/#kit-unpack + * + * @example + * ```ts + * const op = unpack('registry.example.com/org/model:v1', { dir: './model' }); + * setTimeout(() => op.cancel(), 60_000); + * try { + * await op; + * } catch (e) { + * if (e instanceof DOMException && e.name === 'AbortError') { + * console.log('Unpack was cancelled'); + * } + * } + * ``` */ -export async function unpack(path: string, flags?: UnpackFlags): Promise { - const args = [path] - - if (flags) { - args.push(...prepareArgs(flags)) - } - - await runCommand('unpack', args); -} \ No newline at end of file +export function unpack(path: string, flags?: UnpackFlags): CancellablePromise { + return cancellable((signal) => { + const args = [path, ...(flags ? prepareArgs(flags) : [])]; + return runCommand('unpack', args, undefined, { signal }).then(() => {}); + }); +} diff --git a/src/core/__tests__/exec.spec.ts b/src/core/__tests__/exec.spec.ts index bd0a6e8..ab6af90 100644 --- a/src/core/__tests__/exec.spec.ts +++ b/src/core/__tests__/exec.spec.ts @@ -1,6 +1,6 @@ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { spawn } from 'child_process'; -import { runCommand, prepareArgs, parseTableOutput } from '../exec'; +import { runCommand, prepareArgs, parseTableOutput, cancellable } from '../exec'; vi.mock('child_process'); @@ -138,6 +138,128 @@ describe('exec', () => { }); }); + describe('AbortSignal support', () => { + it('should reject immediately without spawning when signal is already aborted', async () => { + const ac = new AbortController(); + ac.abort(); + + const promise = runCommand('version', [], undefined, { signal: ac.signal }); + + await expect(promise).rejects.toMatchObject({ name: 'AbortError' }); + expect(mockSpawn).not.toHaveBeenCalled(); + }); + + it('should kill the child process and reject when signal is aborted mid-run', async () => { + const mockChild = { + stdout: { on: vi.fn() }, + stderr: { on: vi.fn() }, + on: vi.fn(), + kill: vi.fn(), + }; + + mockSpawn.mockReturnValue(mockChild as any); + + const ac = new AbortController(); + const promise = runCommand('version', [], undefined, { signal: ac.signal }); + + ac.abort(); + + const closeCallback = mockChild.on.mock.calls.find(call => call[0] === 'close')?.[1]; + closeCallback?.(null); + + await expect(promise).rejects.toMatchObject({ name: 'AbortError' }); + expect(mockChild.kill).toHaveBeenCalledWith('SIGTERM'); + }); + + it('should clean up the abort listener after the process closes normally', async () => { + const mockChild = { + stdout: { on: vi.fn() }, + stderr: { on: vi.fn() }, + on: vi.fn(), + kill: vi.fn(), + }; + + mockSpawn.mockReturnValue(mockChild as any); + + const ac = new AbortController(); + const removeEventListenerSpy = vi.spyOn(ac.signal, 'removeEventListener'); + + const promise = runCommand('version', [], undefined, { signal: ac.signal }); + + const closeCallback = mockChild.on.mock.calls.find(call => call[0] === 'close')?.[1]; + closeCallback?.(0); + + await promise; + + expect(removeEventListenerSpy).toHaveBeenCalled(); + }); + + it('should resolve normally when signal is provided but never aborted', async () => { + const mockChild = { + stdout: { on: vi.fn() }, + stderr: { on: vi.fn() }, + on: vi.fn(), + kill: vi.fn(), + }; + + mockSpawn.mockReturnValue(mockChild as any); + + const ac = new AbortController(); + const promise = runCommand('version', [], undefined, { signal: ac.signal }); + + const stdoutDataCallback = mockChild.stdout.on.mock.calls.find(call => call[0] === 'data')?.[1]; + const closeCallback = mockChild.on.mock.calls.find(call => call[0] === 'close')?.[1]; + + stdoutDataCallback?.('v1.0.0\n'); + closeCallback?.(0); + + await expect(promise).resolves.toEqual({ stdout: 'v1.0.0', stderr: '', exitCode: 0 }); + expect(mockChild.kill).not.toHaveBeenCalled(); + }); + }); + + describe('cancellable', () => { + it('should return a Promise with a cancel method', () => { + const op = cancellable(() => Promise.resolve()); + expect(op).toBeInstanceOf(Promise); + expect(typeof op.cancel).toBe('function'); + }); + + it('should resolve the underlying promise', async () => { + const op = cancellable(() => Promise.resolve(42)); + await expect(op).resolves.toBe(42); + }); + + it('should reject when the underlying promise rejects', async () => { + const op = cancellable(() => Promise.reject(new Error('fail'))); + await expect(op).rejects.toThrow('fail'); + }); + + it('should abort the signal when cancel is called', () => { + let capturedSignal: AbortSignal | undefined; + + const op = cancellable((signal) => { + capturedSignal = signal; + return new Promise(() => {}); + }); + + expect(capturedSignal?.aborted).toBe(false); + op.cancel(); + expect(capturedSignal?.aborted).toBe(true); + }); + + it('should pass an AbortSignal to the factory', () => { + let capturedSignal: AbortSignal | undefined; + + cancellable((signal) => { + capturedSignal = signal; + return Promise.resolve(); + }); + + expect(capturedSignal).toBeInstanceOf(AbortSignal); + }); + }); + describe('prepareArgs', () => { it('should handle boolean flags', () => { const options = { force: true, quiet: false, verbose: true }; diff --git a/src/core/exec.ts b/src/core/exec.ts index b5a72c1..8e68c93 100644 --- a/src/core/exec.ts +++ b/src/core/exec.ts @@ -1,5 +1,5 @@ import { spawn } from 'child_process'; -import type { KitCommand, ExecResult } from '../types/kitops.js'; +import type { KitCommand, ExecResult, CancellablePromise } from '../types/kitops.js'; type ParsedTableResult = { [key: string]: string } @@ -9,6 +9,7 @@ type ExecOptions = { cwd?: string, env?: Record, stdio?: StdIO | StdIO[], + signal?: AbortSignal, } /** @@ -19,20 +20,50 @@ function getKitCli() { return process.env.KITOPS_CLI_PATH || 'kit'; } +/** + * Returns a `DOMException` named `'AbortError'`, passing the value through + * unchanged when it already is one. Any other reason is attached as `cause`. + */ +function toAbortError(reason: unknown): DOMException { + if (reason instanceof DOMException && reason.name === 'AbortError') { + return reason + } + + const err = new DOMException('The operation was aborted', 'AbortError'); + if (reason !== undefined) { + Object.assign(err, { cause: reason }); + } + + return err; +} + /** * Spawns the kit CLI and runs a single command, returning its captured output. * - * Rejects if the process emits an 'error' event (e.g. binary not found) or - * exits with a non-zero code. The rejection value is a plain string that - * includes the exit code and stderr text so callers don't have to reconstruct it. + * Rejects in three ways: + * - `string` — spawn error (e.g. binary not found) or non-zero exit code; includes + * the exit code and stderr text. + * - `DOMException` named `'AbortError'` — the signal was aborted; a non-AbortError + * signal reason is available as `err.cause`. + * + * When `options.signal` is provided and aborted, the child process is sent `SIGTERM` + * and the promise always rejects with a `DOMException` named `'AbortError'`. If the + * signal carries a non-AbortError reason it is attached as `err.cause`. If the signal + * is already aborted before spawning, the process is never started. * * @param stdin - Data to write to stdin before closing it (e.g., a password) * @param options.cwd - Working directory for the spawned process (default: process.cwd()) * @param options.env - Extra environment variables merged on top of process.env * @param options.stdio - stdio configuration for the spawned process (default: 'pipe' to capture output) + * @param options.signal - AbortSignal that kills the child process when fired */ export function runCommand(command: KitCommand, args: string[] = [], stdin?: string, options: ExecOptions = {}): Promise { return new Promise((resolve, reject) => { + if (options.signal?.aborted) { + reject(toAbortError(options.signal.reason)); + return; + } + const fullArgs = [command, ...args]; const child = spawn(getKitCli(), fullArgs, { cwd: options.cwd || process.cwd(), @@ -42,6 +73,20 @@ export function runCommand(command: KitCommand, args: string[] = [], stdin?: str let stdout = ''; let stderr = ''; + let aborted = false; + + const onAbort = () => { + aborted = true; + child.kill('SIGTERM'); + }; + + const cleanup = () => { + options.signal?.removeEventListener('abort', onAbort); + }; + + if (options.signal) { + options.signal.addEventListener('abort', onAbort, { once: true }); + } if (stdin && child.stdin) { child.stdin.write(stdin); @@ -61,10 +106,18 @@ export function runCommand(command: KitCommand, args: string[] = [], stdin?: str } child.on('error', (error) => { + cleanup(); reject(`Failed to execute kit command: ${error.message}`); }); child.on('close', (code) => { + cleanup(); + + if (aborted) { + reject(toAbortError(options.signal?.reason)); + return; + } + const result: ExecResult = { stdout: stdout.trim(), stderr: stderr.trim(), @@ -80,6 +133,26 @@ export function runCommand(command: KitCommand, args: string[] = [], stdin?: str }); } +/** + * Wraps an async factory in a {@link CancellablePromise}. + * + * The library creates and owns the `AbortController` internally — callers never need to + * set one up. Call `.cancel()` on the returned promise to abort the operation at any time. + * + * Falls back gracefully to a no-op `.cancel()` in environments where `AbortController` + * is unavailable, so the operation still runs to completion safely. + * + * @param fn - Factory that receives an optional `AbortSignal` and returns a `Promise`. + */ +export function cancellable(fn: (signal: AbortSignal | undefined) => Promise): CancellablePromise { + if (typeof AbortController === 'undefined') { + return Object.assign(fn(undefined), { cancel: () => {} }); + } + + const ac = new AbortController(); + return Object.assign(fn(ac.signal), { cancel: () => ac.abort() }); +} + /** * Converts a flags object into a flat array of CLI arguments. * camelCase keys are converted to kebab-case. diff --git a/src/index.ts b/src/index.ts index 56f3872..93eebdf 100644 --- a/src/index.ts +++ b/src/index.ts @@ -56,7 +56,8 @@ export type { export type { ModelKit, ExecResult, - KitCommand + KitCommand, + CancellablePromise } from './types/kitops.js'; export type { diff --git a/src/types/kitops.ts b/src/types/kitops.ts index b85f85a..66efbd9 100644 --- a/src/types/kitops.ts +++ b/src/types/kitops.ts @@ -53,4 +53,32 @@ export type ExecResult = { stdout: string; stderr: string; exitCode: number; +}; + +/** + * A `Promise` that exposes a `cancel()` method to abort the underlying `kit` process. + * + * Calling `cancel()` sends `SIGTERM` to the child process and rejects the promise with a + * `DOMException` named `'AbortError'`. In environments where `AbortController` is unavailable, + * `cancel()` is a no-op and the operation runs to completion normally. + * + * Because `CancellablePromise` extends `Promise`, existing call sites using `await` + * require no changes. Hold the reference before awaiting to gain access to `cancel()`. + * + * @example + * ```ts + * const op = pull('registry.example.com/org/model:v1'); + * setTimeout(() => op.cancel(), 30_000); + * try { + * await op; + * } catch (e) { + * if (e instanceof DOMException && e.name === 'AbortError') { + * console.log('Pull was cancelled'); + * } + * } + * ``` + */ +export type CancellablePromise = Promise & { + /** Kills the underlying `kit` process and rejects the promise with an `AbortError`. */ + cancel: () => void; }; \ No newline at end of file