255 lines
16 KiB
JavaScript
255 lines
16 KiB
JavaScript
import test from 'node:test';
|
||
import assert from 'node:assert/strict';
|
||
import { createServer } from 'node:http';
|
||
import { mkdtemp, readFile, writeFile, rm } from 'node:fs/promises';
|
||
import { tmpdir } from 'node:os';
|
||
import { join, dirname, basename } from 'node:path';
|
||
|
||
// Match DSH credentials credentialKey/parseCredentialKey admission rules.
|
||
// Validate at every service boundary AND on reload, unlike the original permissive mock.
|
||
function validateCredentialName(key) {
|
||
assert.match(key, /^[a-z][a-z0-9-]*\/[a-z][a-z0-9-]*$/);
|
||
}
|
||
|
||
const PNG = Buffer.from('iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+jVv8AAAAASUVORK5CYII=', 'base64');
|
||
|
||
test('native host integration using only a local image API', async t => {
|
||
const home = await mkdtemp(join(tmpdir(), 'dsh-imagegen-host-test-'));
|
||
const previousHome = process.env.DSH_HOME;
|
||
process.env.DSH_HOME = home;
|
||
t.after(async () => {
|
||
if (previousHome === undefined) delete process.env.DSH_HOME;
|
||
else process.env.DSH_HOME = previousHome;
|
||
assert.equal(dirname(home), tmpdir());
|
||
assert.match(basename(home), /^dsh-imagegen-host-test-/);
|
||
await rm(home, { recursive: true, force: true });
|
||
});
|
||
const requests = [];
|
||
const server = createServer(async (request, response) => {
|
||
try {
|
||
const chunks = [];
|
||
for await (const chunk of request) chunks.push(chunk);
|
||
const raw = Buffer.concat(chunks);
|
||
const body = request.url === '/v1/images/edits'
|
||
? await new Response(raw, { headers: { 'Content-Type': request.headers['content-type'] } }).formData()
|
||
: raw.length ? JSON.parse(raw.toString()) : undefined;
|
||
requests.push({ path: request.url, authorization: request.headers.authorization, body });
|
||
response.setHeader('Content-Type', 'application/json');
|
||
if (request.url === '/v1/models') response.end(JSON.stringify({ data: [{ id: 'test-image-model' }] }));
|
||
else if (['/v1/images/generations', '/v1/images/edits'].includes(request.url)) response.end(JSON.stringify({ data: [{ b64_json: PNG.toString('base64'), revised_prompt: 'Local test image' }] }));
|
||
else { response.statusCode = 404; response.end('{}'); }
|
||
} catch { response.statusCode = 500; response.end('{}'); }
|
||
});
|
||
await new Promise((resolve, reject) => { server.once('error', reject); server.listen(0, '127.0.0.1', resolve); });
|
||
t.after(() => new Promise((resolve, reject) => {
|
||
server.close(error => error ? reject(error) : resolve());
|
||
server.closeAllConnections();
|
||
}));
|
||
const baseUrl = `http://127.0.0.1:${server.address().port}/v1`;
|
||
const records = new Map(), routes = new Map(), tools = new Map(), skills = [], attachments = [], listeners = new Map();
|
||
const ctx = {
|
||
on(event, listener) { listeners.set(event, listener); },
|
||
connection: { fetch: { register(route) {
|
||
assert.equal(route.requestBody, route.path === '/api/imagegen/edit' ? 'streaming' : 'buffered');
|
||
assert.ok(!routes.has(route.path));
|
||
routes.set(route.path, route);
|
||
return async () => routes.delete(route.path);
|
||
} } },
|
||
tools: { register(tool) { tools.set(tool.name, tool); return () => tools.delete(tool.name); } },
|
||
credentials: {
|
||
async readRecord(key) { validateCredentialName(key); return records.get(key); },
|
||
async modifyRecord(key, change) { validateCredentialName(key); const value = await change(records.get(key)); records.set(key, value); return value; },
|
||
async deleteRecord(key) { validateCredentialName(key); records.delete(key); },
|
||
},
|
||
attachments: { async validateImage(image) { assert.deepEqual(Buffer.from(image.data), PNG); }, async saveImage(image) {
|
||
assert.deepEqual(Buffer.from(image.data), PNG);
|
||
assert.equal(image.mediaType, 'image/png');
|
||
const ref = { attachmentId: `test-attachment-${attachments.length}`, mediaType: image.mediaType, bytes: image.data.length, width: 1, height: 1, name: image.name };
|
||
attachments.push(ref);
|
||
return ref;
|
||
} },
|
||
skills: { register(skill) { skills.push(skill); return () => {}; } },
|
||
};
|
||
const host = await import(`../lib/index.mjs?host-test=${Date.now()}`);
|
||
await host.apply(ctx);
|
||
t.after(() => listeners.get('dispose')?.());
|
||
const call = async (path, method = 'GET', body) => {
|
||
const url = new URL(path, 'http://localhost'), route = routes.get(url.pathname);
|
||
assert.ok(route, `route registered: ${url.pathname}`);
|
||
assert.ok(route.methods.includes(method));
|
||
return route.fetch(new Request(url, { method, ...(body === undefined ? {} : { body: JSON.stringify(body), headers: { 'Content-Type': 'application/json' } }) }));
|
||
};
|
||
const settings = { baseUrl, models: ['test-image-model'], defaultModel: 'test-image-model' };
|
||
const key = 'test-only-key-never-a-real-secret';
|
||
|
||
await t.test('registers dependencies, routes, tools and skill', () => {
|
||
assert.deepEqual(host.inject, ['connection', 'tools', 'credentials', 'attachments', 'skills']);
|
||
assert.deepEqual([...routes.keys()].sort(), ['/api/imagegen/edit', '/api/imagegen/file', '/api/imagegen/generate', '/api/imagegen/models', '/api/imagegen/settings']);
|
||
assert.deepEqual([...tools.keys()].sort(), ['imagegen_edit', 'imagegen_generate', 'imagegen_status']);
|
||
assert.equal(skills.length, 1);
|
||
assert.equal(skills[0].name, 'dsh-imagegen');
|
||
assert.ok(skills[0].content.length > 0);
|
||
assert.equal(typeof listeners.get('dispose'), 'function');
|
||
});
|
||
await t.test('stores keys in credentials, not settings JSON or returned values', async () => {
|
||
const response = await call('/api/imagegen/settings', 'POST', { ...settings, apiKey: key });
|
||
assert.equal(response.status, 200);
|
||
const value = await response.json();
|
||
assert.equal(value.configured, true);
|
||
assert.equal(value.apiKey, undefined);
|
||
assert.ok(!JSON.stringify(value).includes(key));
|
||
assert.equal(records.size, 1);
|
||
const [credentialName, credential] = [...records][0];
|
||
assert.match(credentialName, /^imagegen-local\/endpoint-[0-9a-f]{64}$/);
|
||
assert.deepEqual(credential, { kind: 'api-key', key });
|
||
const disk = await readFile(join(home, 'imagegen', 'settings.json'), 'utf8');
|
||
assert.deepEqual(JSON.parse(disk), settings);
|
||
assert.ok(!disk.includes(key));
|
||
const status = await tools.get('imagegen_status').execute();
|
||
assert.equal(status.configured, true);
|
||
assert.ok(!JSON.stringify(status).includes(key));
|
||
assert.ok(!(await (await call('/api/imagegen/settings')).text()).includes(key));
|
||
});
|
||
await t.test('saved credentials survive strict reload and host reinitialization', async () => {
|
||
const persisted = JSON.stringify({ version: 1, records: Object.fromEntries(records) });
|
||
records.clear();
|
||
for (const [address, value] of Object.entries(JSON.parse(persisted).records)) {
|
||
validateCredentialName(address);
|
||
records.set(address, value);
|
||
}
|
||
listeners.get('dispose')();
|
||
routes.clear(); tools.clear();
|
||
await host.apply(ctx);
|
||
assert.equal((await tools.get('imagegen_status').execute()).configured, true);
|
||
});
|
||
await t.test('lists models using the local authenticated API', async () => {
|
||
const response = await call('/api/imagegen/models', 'POST');
|
||
assert.equal(response.status, 200);
|
||
assert.deepEqual(await response.json(), { models: ['test-image-model'] });
|
||
assert.equal(requests.at(-1).authorization, `Bearer ${key}`);
|
||
});
|
||
await t.test('a different Base URL cannot reuse the original credential', async () => {
|
||
const response = await call('/api/imagegen/settings', 'POST', { ...settings, baseUrl: `${baseUrl}/other-provider` });
|
||
assert.equal(response.status, 200);
|
||
assert.equal((await response.json()).configured, false);
|
||
const before = requests.length;
|
||
const rejected = await call('/api/imagegen/generate', 'POST', { model: 'test-image-model', prompt: 'A local test' });
|
||
assert.equal(rejected.status, 400);
|
||
assert.equal(requests.length, before);
|
||
assert.equal(records.size, 1);
|
||
const restored = await call('/api/imagegen/settings', 'POST', settings);
|
||
assert.equal((await restored.json()).configured, true);
|
||
});
|
||
await t.test('requires explicit model despite a configured default', async () => {
|
||
const before = requests.length;
|
||
const response = await call('/api/imagegen/generate', 'POST', { prompt: 'No model' });
|
||
assert.equal(response.status, 400);
|
||
assert.match((await response.json()).error, /model/i);
|
||
await assert.rejects(tools.get('imagegen_generate').execute({ prompt: 'No model' }, { signal: new AbortController().signal }), /model/i);
|
||
assert.equal(requests.length, before);
|
||
});
|
||
await t.test('tool generation saves original bytes and renders an image attachment', async () => {
|
||
const tool = tools.get('imagegen_generate');
|
||
const args = { model: 'test-image-model', prompt: 'A local test', n: 1 };
|
||
const value = await tool.execute(args, { signal: new AbortController().signal });
|
||
assert.equal(value.model, args.model);
|
||
assert.equal(value.images.length, 1);
|
||
const image = value.images[0];
|
||
assert.equal(dirname(image.path), join(home, 'imagegen', 'outputs'));
|
||
assert.deepEqual(await readFile(image.path), PNG);
|
||
assert.equal(image.revisedPrompt, 'Local test image');
|
||
assert.deepEqual(image.attachment, attachments[0]);
|
||
const content = tool.output.render(args, value);
|
||
assert.equal(content[0].type, 'text');
|
||
assert.deepEqual(content[1], { type: 'image', attachment: attachments[0] });
|
||
assert.ok(!JSON.stringify(content).includes(key));
|
||
assert.deepEqual(requests.at(-1).body, args);
|
||
assert.equal(requests.at(-1).authorization, `Bearer ${key}`);
|
||
const file = await call(image.url);
|
||
assert.equal(file.status, 200);
|
||
assert.equal(file.headers.get('Content-Type'), 'image/png');
|
||
assert.equal(file.headers.get('X-Content-Type-Options'), 'nosniff');
|
||
assert.deepEqual(Buffer.from(await file.arrayBuffer()), PNG);
|
||
});
|
||
await t.test('HTTP generation returns a saved image without credentials', async () => {
|
||
const response = await call('/api/imagegen/generate', 'POST', { model: 'test-image-model', prompt: 'HTTP local test' });
|
||
assert.equal(response.status, 200);
|
||
const value = await response.json();
|
||
assert.equal(value.images.length, 1);
|
||
assert.deepEqual(await readFile(value.images[0].path), PNG);
|
||
assert.ok(!JSON.stringify(value).includes(key));
|
||
});
|
||
await t.test('edit tool uploads original bytes and saves separate output', async () => {
|
||
const inputPath = join(home, 'input.png');
|
||
await writeFile(inputPath, PNG);
|
||
const value = await tools.get('imagegen_edit').execute({ image_path: inputPath, model: 'test-image-model', prompt: 'Change background only' }, { signal: new AbortController().signal });
|
||
const request = requests.at(-1);
|
||
assert.equal(request.path, '/v1/images/edits');
|
||
assert.equal(request.authorization, `Bearer ${key}`);
|
||
assert.equal(request.body.get('model'), 'test-image-model');
|
||
assert.deepEqual(Buffer.from(await request.body.get('image').arrayBuffer()), PNG);
|
||
assert.deepEqual(await readFile(inputPath), PNG);
|
||
assert.notEqual(value.images[0].path, inputPath);
|
||
assert.deepEqual(await readFile(value.images[0].path), PNG);
|
||
assert.equal(tools.get('imagegen_edit').output.render({}, value)[1].type, 'image');
|
||
assert.ok(!JSON.stringify(value).includes(key));
|
||
const before = requests.length;
|
||
for (const image_path of ['relative.png', join(home, 'missing.png')]) {
|
||
await assert.rejects(tools.get('imagegen_edit').execute({ image_path, model: 'test-image-model', prompt: 'edit' }, { signal: new AbortController().signal }));
|
||
}
|
||
await assert.rejects(tools.get('imagegen_edit').execute({ image_path: inputPath, prompt: 'edit' }, { signal: new AbortController().signal }), /model/);
|
||
assert.equal(requests.length, before);
|
||
});
|
||
await t.test('multipart edit route accepts one file and rejects invalid uploads before API calls', async () => {
|
||
const send = form => routes.get('/api/imagegen/edit').fetch(new Request('http://localhost/api/imagegen/edit', { method: 'POST', body: form }));
|
||
const make = () => { const form = new FormData(); form.set('input', JSON.stringify({ model: 'test-image-model', prompt: 'Edit via browser' })); form.set('image', new Blob([PNG], { type: 'image/png' }), 'input.png'); return form; };
|
||
const result = await send(make());
|
||
assert.equal(result.status, 200);
|
||
assert.deepEqual(await readFile((await result.json()).images[0].path), PNG);
|
||
const before = requests.length;
|
||
const missing = make(); missing.delete('image');
|
||
const duplicate = make(); duplicate.append('image', new Blob([PNG]), 'second.png');
|
||
const invalid = make(); invalid.set('image', new Blob(['not an image']), 'fake.png');
|
||
const unsupported = make(); unsupported.set('mask', new Blob([PNG]), 'mask.png');
|
||
const invalidJson = make(); invalidJson.set('input', '{');
|
||
const noModel = make(); noModel.set('input', JSON.stringify({ prompt: 'edit' }));
|
||
for (const form of [missing, duplicate, invalid, unsupported, invalidJson, noModel]) assert.equal((await send(form)).status, 400);
|
||
assert.equal((await call('/api/imagegen/edit', 'POST', { image_path: join(home, 'input.png') })).status, 400);
|
||
const tooLarge = new Request('http://localhost/api/imagegen/edit', { method: 'POST', headers: { 'Content-Length': String(22 * 1024 * 1024) }, body: 'x' });
|
||
assert.equal((await routes.get('/api/imagegen/edit').fetch(tooLarge)).status, 400);
|
||
// Model bytes arriving from the browser transport; cancelling Node's own
|
||
// FormData encoder mid-Blob triggers an unrelated undici closed-stream bug.
|
||
const oversized = new Request('http://localhost/api/imagegen/edit', { method: 'POST', headers: { 'Content-Type': 'multipart/form-data; boundary=test' }, body: Buffer.alloc(22 * 1024 * 1024) });
|
||
assert.equal((await routes.get('/api/imagegen/edit').fetch(oversized)).status, 400, 'streamed body limit applies even without Content-Length');
|
||
const decode = ctx.attachments.validateImage;
|
||
ctx.attachments.validateImage = async () => { throw new Error('Image decode refused'); };
|
||
try { assert.equal((await send(make())).status, 400); } finally { ctx.attachments.validateImage = decode; }
|
||
assert.equal(requests.length, before);
|
||
});
|
||
await t.test('all digest prefixes are valid and clearing only removes the selected endpoint', async () => {
|
||
const originalRecords = new Map(records);
|
||
const prefixes = new Set();
|
||
for (let i = 0; prefixes.size < 16 && i < 512; i++) {
|
||
const endpoint = { ...settings, baseUrl: `https://example.com/v1/provider-${i}` };
|
||
const response = await call('/api/imagegen/settings', 'POST', { ...endpoint, apiKey: 'test-only-temporary-key' });
|
||
assert.equal(response.status, 200);
|
||
assert.equal((await response.json()).configured, true);
|
||
const address = [...records.keys()].find(item => !originalRecords.has(item));
|
||
validateCredentialName(address);
|
||
prefixes.add(address.split('endpoint-')[1][0]);
|
||
assert.equal((await (await call('/api/imagegen/settings', 'POST', endpoint)).json()).configured, true);
|
||
const cleared = await call('/api/imagegen/settings', 'POST', { ...endpoint, clearKey: true });
|
||
assert.equal(cleared.status, 200);
|
||
assert.equal((await cleared.json()).configured, false);
|
||
assert.deepEqual(records, originalRecords);
|
||
}
|
||
assert.equal(prefixes.size, 16, 'exercise every hexadecimal leading character, including 0–9');
|
||
await call('/api/imagegen/settings', 'POST', settings);
|
||
});
|
||
await t.test('file traversal, absolute paths and missing files return 404', async () => {
|
||
for (const filename of ['../settings.json', '..\\settings.json', join(home, 'imagegen', 'settings.json'), '/etc/passwd', '', '123-00000000-0000-0000-0000-000000000000.png']) {
|
||
assert.equal((await call(`/api/imagegen/file?name=${encodeURIComponent(filename)}`)).status, 404);
|
||
}
|
||
});
|
||
});
|