aaronrene / knowtation public
app-runtime.mjs
152 lines 5.0 KB
Raw
sha256:bd1eb24434d281fe65142038c1bf217493421f5a86880c22dff7e6037d2d4700 feat(companion): add production macOS distribution path Human minor ⚠ breaking 24 days ago
1 import crypto from 'node:crypto';
2 import fs from 'node:fs/promises';
3 import os from 'node:os';
4 import path from 'node:path';
5 import { fork } from 'node:child_process';
6
7 import { createCompanionInferenceListener } from '../../lib/companion-inference-listener.mjs';
8
9 const MAX_BODY_BYTES = 512 * 1024;
10
11 function waitForWorker(worker) {
12 return new Promise((resolve, reject) => {
13 const timer = setTimeout(() => reject(new Error('PRIVATE_MODEL_UNAVAILABLE')), 30_000);
14 timer.unref?.();
15 const onMessage = (message) => {
16 if (message?.ready !== true) return;
17 clearTimeout(timer);
18 worker.off('message', onMessage);
19 worker.off('exit', onExit);
20 resolve();
21 };
22 const onExit = () => {
23 clearTimeout(timer);
24 worker.off('message', onMessage);
25 reject(new Error('PRIVATE_MODEL_UNAVAILABLE'));
26 };
27 worker.on('message', onMessage);
28 worker.once('exit', onExit);
29 });
30 }
31
32 function createWorkerClient(worker) {
33 const pending = new Map();
34 worker.on('message', (message) => {
35 if (!message?.id || !pending.has(message.id)) return;
36 const entry = pending.get(message.id);
37 pending.delete(message.id);
38 clearTimeout(entry.timer);
39 if (message.ok) entry.resolve(message.result);
40 else entry.reject(new Error(message.error ?? 'PRIVATE_MODEL_UNAVAILABLE'));
41 });
42 worker.once('exit', () => {
43 for (const entry of pending.values()) {
44 clearTimeout(entry.timer);
45 entry.reject(new Error('PRIVATE_MODEL_UNAVAILABLE'));
46 }
47 pending.clear();
48 });
49 return {
50 embed(texts) {
51 const id = crypto.randomUUID();
52 return new Promise((resolve, reject) => {
53 const timer = setTimeout(() => {
54 pending.delete(id);
55 reject(new Error('PRIVATE_MODEL_UNAVAILABLE'));
56 }, 60_000);
57 timer.unref?.();
58 pending.set(id, { resolve, reject, timer });
59 worker.send({ id, command: 'embed', arguments: { texts } });
60 });
61 },
62 };
63 }
64
65 async function readJsonBody(request) {
66 const chunks = [];
67 let bytes = 0;
68 for await (const chunk of request) {
69 bytes += chunk.length;
70 if (bytes > MAX_BODY_BYTES) throw new Error('PRIVATE_INFERENCE_INPUT');
71 chunks.push(chunk);
72 }
73 return JSON.parse(Buffer.concat(chunks).toString('utf8'));
74 }
75
76 export async function startAppRuntime({
77 custody,
78 runtimeRoot,
79 spawnWorker = fork,
80 listenerFactory = createCompanionInferenceListener,
81 stateDirectory = path.join(os.homedir(), 'Library/Application Support/Knowtation'),
82 }) {
83 const loopbackToken = crypto.randomBytes(32).toString('base64url');
84 const tokenAccount = 'knowtation.companion.loopbackToken';
85 const workerPath = path.join(runtimeRoot, 'companion/runtime/inference-worker.mjs');
86 const stateFile = path.join(stateDirectory, 'companion.json');
87 let worker;
88 let listener;
89 let closed = false;
90
91 async function cleanup() {
92 if (closed) return;
93 closed = true;
94 const tasks = [];
95 if (listener) tasks.push(Promise.resolve().then(() => listener.close()));
96 if (worker) worker.kill('SIGTERM');
97 tasks.push(Promise.resolve().then(() => custody.delete(tokenAccount)));
98 tasks.push(fs.rm(stateFile, { force: true }));
99 const results = await Promise.allSettled(tasks);
100 const failure = results.find((result) => result.status === 'rejected');
101 if (failure) throw failure.reason;
102 }
103
104 try {
105 await custody.set(tokenAccount, loopbackToken);
106 worker = spawnWorker(workerPath, [runtimeRoot], {
107 execPath: process.execPath,
108 execArgv: ['--permission', `--allow-fs-read=${runtimeRoot}`],
109 env: {
110 PATH: '/usr/bin:/bin:/usr/sbin:/sbin',
111 LANG: process.env.LANG ?? 'en_US.UTF-8',
112 TMPDIR: process.env.TMPDIR ?? os.tmpdir(),
113 },
114 stdio: ['ignore', 'ignore', 'ignore', 'ipc'],
115 });
116 await waitForWorker(worker);
117 const client = createWorkerClient(worker);
118
119 listener = listenerFactory({
120 expectedToken: loopbackToken,
121 async runtimeRequest(request, response) {
122 response.setHeader('Content-Type', 'application/json');
123 response.setHeader('Cache-Control', 'no-store');
124 if (request.method === 'GET') {
125 response.statusCode = 200;
126 response.end(JSON.stringify({ ok: true, model: 'all-MiniLM-L6-v2', dimensions: 384 }));
127 return;
128 }
129 const body = await readJsonBody(request);
130 const result = await client.embed(body?.texts);
131 response.statusCode = 200;
132 response.end(JSON.stringify({ ok: true, ...result }));
133 },
134 });
135 const endpoint = await listener.start();
136
137 await fs.mkdir(stateDirectory, { recursive: true, mode: 0o700 });
138 await fs.writeFile(stateFile, `${JSON.stringify({
139 schema: 1,
140 host: endpoint.host,
141 port: endpoint.port,
142 pid: process.pid,
143 version: process.env.KNOWTATION_RELEASE_VERSION,
144 build: Number(process.env.KNOWTATION_RELEASE_BUILD),
145 })}\n`, { mode: 0o600 });
146
147 return Object.freeze({ endpoint, close: cleanup });
148 } catch (error) {
149 await cleanup().catch(() => {});
150 throw error;
151 }
152 }
File History 1 commit
sha256:bd1eb24434d281fe65142038c1bf217493421f5a86880c22dff7e6037d2d4700 feat(companion): add production macOS distribution path Human minor ⚠ 24 days ago