omnikey-cli 1.0.35 → 1.0.36

package/backend-dist/agent/agentServer.js

@@ -171,9 +171,6 @@ async function runToolLoop(initialResult, session, sessionId, send, log, tools,
     return result;
 }
 const aiModel = (0, ai_client_1.getDefaultModel)(config_1.config.aiProvider, 'smart');
-// In-memory cache: sessionId -> live SessionState. Hydrated from DB on first
-// access and written back after each turn so restarts resume correctly.
-const sessionMessages = new Map();
 const MAX_TURNS = 20;
 // ─── DB helpers ───────────────────────────────────────────────────────────────
 async function persistSessionToDB(sessionId, state) {
@@ -216,22 +213,7 @@ async function enforceSessionCap(subscriptionId, logger) {
     }
 }
 async function getOrCreateSession(sessionId, subscription, platform, log, isCronJob = false) {
-    // 1. Return the live in-memory entry if already loaded this process lifetime.
-    const existing = sessionMessages.get(sessionId);
-    if (existing) {
-        log.debug('Reusing existing agent session (in-memory)', {
-            sessionId,
-            subscriptionId: existing.subscription.id,
-            turns: existing.turns,
-        });
-        return {
-            sessionState: existing,
-            hasStoredPrompt: existing.history
-                .filter((h) => h.role === 'user')
-                .some((h) => typeof h.content === 'string' && h.content.includes('<stored_instructions>')),
-        };
-    }
-    // 2. Try to resume from a persisted DB record.
+    // 1. Try to resume from a persisted DB record.
     try {
         const dbSession = await agentSession_1.AgentSession.findOne({
             where: { id: sessionId, subscriptionId: subscription.id },
@@ -243,7 +225,6 @@ async function getOrCreateSession(sessionId, subscription, platform, log, isCron
                 history,
                 turns: dbSession.turns,
             };
-            sessionMessages.set(sessionId, entry);
             log.info('Resumed agent session from DB', {
                 sessionId,
                 subscriptionId: subscription.id,
@@ -263,7 +244,7 @@ async function getOrCreateSession(sessionId, subscription, platform, log, isCron
             error: err,
         });
     }
-    // 3. Create a brand-new session in-memory and persist it to the DB.
+    // 2. Create a brand-new session and persist it to the DB.
     const prompt = await (0, featureRoutes_1.getPromptForCommand)(log, 'task', subscription).catch((err) => {
         log.error('Failed to get system prompt for new agent session', { error: err });
         return '';
@@ -293,17 +274,19 @@ ${prompt}
         ],
         turns: 0,
     };
-    sessionMessages.set(sessionId, entry);
     // Persist immediately so that GET /sessions picks it up right away.
     try {
-        await agentSession_1.AgentSession.create({
-            id: sessionId,
-            subscriptionId: subscription.id,
-            title: 'New Session',
-            platform: platform ?? null,
-            historyJson: JSON.stringify(entry.history),
-            turns: 0,
-            lastActiveAt: new Date(),
+        await agentSession_1.AgentSession.findOrCreate({
+            where: { id: sessionId, subscriptionId: subscription.id },
+            defaults: {
+                id: sessionId,
+                subscriptionId: subscription.id,
+                title: 'New Session',
+                platform: platform ?? null,
+                historyJson: JSON.stringify(entry.history),
+                turns: 0,
+                lastActiveAt: new Date(),
+            },
         });
         // Prune oldest sessions after each creation so the cap is always respected.
         void enforceSessionCap(subscription.id, log);
@@ -321,7 +304,7 @@ ${prompt}
         hasStoredPrompt: !!prompt,
     };
 }
-async function runAgentTurn(sessionId, subscription, clientMessage, send, log, options) {
+async function runAgentTurnInternal(sessionId, subscription, clientMessage, send, log, options) {
     const { sessionState: session, hasStoredPrompt } = await getOrCreateSession(sessionId, subscription, clientMessage.platform, log, options?.isCronJob);
     // Count this call as one agent iteration.
     session.turns += 1;
@@ -445,7 +428,6 @@ async function runAgentTurn(sessionId, subscription, clientMessage, send, log, o
             log.warn('Agent LLM returned empty content; sending generic error to client.');
             const errorMessage = 'The agent returned an empty response. Please try again.';
             await persistSessionToDB(sessionId, session);
-            sessionMessages.delete(sessionId);
             (0, utils_1.sendFinalAnswer)(send, sessionId, errorMessage, true);
             return;
         }
@@ -506,7 +488,7 @@ async function runAgentTurn(sessionId, subscription, clientMessage, send, log, o
                             'No plain text. No other format.',
                         ].join('\n'),
                 });
-                await runAgentTurn(sessionId, subscription, {
+                await runAgentTurnInternal(sessionId, subscription, {
                     sender: 'agent',
                     session_id: sessionId,
                     content: '',
@@ -552,7 +534,6 @@ async function runAgentTurn(sessionId, subscription, clientMessage, send, log, o
             });
             (0, utils_1.pushToSessionHistory)(logger_1.logger, session, { role: 'assistant', content });
             await persistSessionToDB(sessionId, session);
-            sessionMessages.delete(sessionId);
             send({
                 session_id: sessionId,
                 sender: 'agent',
@@ -571,7 +552,6 @@ async function runAgentTurn(sessionId, subscription, clientMessage, send, log, o
             });
             (0, utils_1.pushToSessionHistory)(log, session, { role: 'assistant', content });
             await persistSessionToDB(sessionId, session);
-            sessionMessages.delete(sessionId);
             send({
                 session_id: sessionId,
                 sender: 'agent',
@@ -583,7 +563,6 @@ async function runAgentTurn(sessionId, subscription, clientMessage, send, log, o
                 sessionId,
             });
             await persistSessionToDB(sessionId, session);
-            sessionMessages.delete(sessionId);
             (0, utils_1.sendFinalAnswer)(send, sessionId, 'The agent returned an empty response. Please try again.', true);
         }
     }
@@ -591,10 +570,12 @@ async function runAgentTurn(sessionId, subscription, clientMessage, send, log, o
         log.error('Agent LLM call failed', { error: err });
         const errorMessage = 'Agent failed to call language model. Please try again later.';
         await persistSessionToDB(sessionId, session);
-        sessionMessages.delete(sessionId);
         (0, utils_1.sendFinalAnswer)(send, sessionId, errorMessage, true);
     }
 }
+async function runAgentTurn(sessionId, subscription, clientMessage, send, log, options) {
+    await runAgentTurnInternal(sessionId, subscription, clientMessage, send, log, options);
+}
 function attachAgentWebSocketServer(server) {
     const wss = new ws_1.WebSocketServer({ server, path: '/ws/omni-agent' });
     wss.on('connection', (ws, req) => {
@@ -731,8 +712,6 @@ function createAgentRouter() {
                 res.status(404).json({ error: 'Session not found' });
                 return;
             }
-            // Also remove from the in-memory cache if it was loaded.
-            sessionMessages.delete(sessionId);
             res.status(200).json({ deleted: true });
         }
         catch (err) {

package/backend-dist/authMiddleware.js

@@ -10,18 +10,33 @@ const jsonwebtoken_1 = __importDefault(require("jsonwebtoken"));
 const logger_1 = require("./logger");
 const config_1 = require("./config");
 const subscription_1 = require("./models/subscription");
+const SELF_HOSTED_SUBSCRIPTION_ID = 'self-hosted-local-subscription';
 async function selfHostedSubscription() {
     try {
-        let subscription = await subscription_1.Subscription.findOne({ where: { isSelfHosted: true } });
-        if (!subscription) {
-            subscription = await subscription_1.Subscription.create({
+        // Reuse any existing self-hosted record (including legacy IDs) first.
+        const existing = await subscription_1.Subscription.findOne({ where: { isSelfHosted: true } });
+        if (existing)
+            return existing;
+        // Use a deterministic primary key so concurrent first-time requests do not
+        // create duplicate rows.
+        const [subscription, created] = await subscription_1.Subscription.findOrCreate({
+            where: { id: SELF_HOSTED_SUBSCRIPTION_ID },
+            defaults: {
+                id: SELF_HOSTED_SUBSCRIPTION_ID,
                 email: 'local-user@omnikey.ai',
                 licenseKey: 'self-hosted',
                 subscriptionStatus: 'active',
                 isSelfHosted: true,
-            });
+            },
+        });
+        if (created) {
             logger_1.logger.info('Created self-hosted subscription record in database.');
         }
+        // Ensure deterministic row remains flagged for self-hosted mode.
+        if (!subscription.isSelfHosted) {
+            subscription.isSelfHosted = true;
+            await subscription.save();
+        }
         return subscription;
     }
     catch (err) {
@@ -31,37 +46,37 @@ async function selfHostedSubscription() {
 }
 async function authMiddleware(req, res, next) {
     const authHeader = req.headers.authorization;
-    logger_1.logger.defaultMeta = { traceId: (0, crypto_1.randomUUID)() };
+    const requestLogger = logger_1.logger.child({ traceId: (0, crypto_1.randomUUID)() });
+    res.locals.logger = requestLogger;
     if (config_1.config.blockSaas) {
-        logger_1.logger.warn('Blocking SaaS access: rejecting request due to BLOCK_SAAS=true');
+        requestLogger.warn('Blocking SaaS access: rejecting request due to BLOCK_SAAS=true');
         return res.status(403).json({ error: 'SaaS access is blocked.' });
     }
     if (config_1.config.isSelfHosted || !config_1.config.jwtSecret) {
-        logger_1.logger.info('Self-hosted mode: skipping auth middleware.');
+        requestLogger.info('Self-hosted mode: skipping auth middleware.');
         if (config_1.config.isSelfHosted) {
             res.locals.subscription = await selfHostedSubscription();
-            res.locals.logger = logger_1.logger;
         }
         return next();
     }
     if (!authHeader) {
-        logger_1.logger.warn('Missing Authorization header on feature route.');
+        requestLogger.warn('Missing Authorization header on feature route.');
         return res.status(401).json({ error: 'Missing bearer token.' });
     }
     const [scheme, token] = authHeader.split(' ');
     if (scheme !== 'Bearer' || !token) {
-        logger_1.logger.warn('Malformed Authorization header on feature route.');
+        requestLogger.warn('Malformed Authorization header on feature route.');
         return res.status(401).json({ error: 'Invalid authorization header.' });
     }
     try {
         const decoded = jsonwebtoken_1.default.verify(token, config_1.config.jwtSecret);
         const subscription = await subscription_1.Subscription.findByPk(decoded.sid);
         if (!subscription) {
-            logger_1.logger.warn('Subscription not found for JWT.', { sid: decoded.sid });
+            requestLogger.warn('Subscription not found for JWT.', { sid: decoded.sid });
             return res.status(403).json({ error: 'Invalid or expired token.' });
         }
         if (subscription.subscriptionStatus == 'expired') {
-            logger_1.logger.warn('Inactive subscription for JWT.', {
+            requestLogger.warn('Inactive subscription for JWT.', {
                 sid: decoded.sid,
                 status: subscription.subscriptionStatus,
             });
@@ -71,19 +86,18 @@ async function authMiddleware(req, res, next) {
         if (subscription.licenseKeyExpiresAt && subscription.licenseKeyExpiresAt <= now) {
             subscription.subscriptionStatus = 'expired';
             await subscription.save();
-            logger_1.logger.info('Subscription key has expired during activation.', {
+            requestLogger.info('Subscription key has expired during activation.', {
                 subscriptionId: subscription.id,
             });
             return res
                 .status(403)
                 .json({ error: 'Subscription has expired.', subscriptionStatus: 'expired' });
         }
-        res.locals.logger = logger_1.logger;
         res.locals.subscription = subscription;
         next();
     }
     catch (err) {
-        logger_1.logger.warn('Invalid or expired JWT on feature route.', { error: err });
+        requestLogger.warn('Invalid or expired JWT on feature route.', { error: err });
         return res.status(403).json({ error: 'Invalid or expired token.' });
     }
 }

package/backend-dist/bucket-adapter/index.js

@@ -3,9 +3,25 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.getDownloadCounts = getDownloadCounts;
 exports.incrementDownloadCount = incrementDownloadCount;
 const storage_1 = require("@google-cloud/storage");
+const zod_1 = require("zod");
 const logger_1 = require("../logger");
 const config_1 = require("../config");
 const DEFAULT_COUNTS = { macos: 0, windows: 0 };
+const downloadCountsSchema = zod_1.z.object({
+    macos: zod_1.z.number().nonnegative().optional(),
+    windows: zod_1.z.number().nonnegative().optional(),
+});
+function parseDownloadCounts(raw) {
+    const json = JSON.parse(raw);
+    const parsed = downloadCountsSchema.safeParse(json);
+    if (!parsed.success) {
+        return { ...DEFAULT_COUNTS };
+    }
+    return {
+        macos: parsed.data.macos ?? 0,
+        windows: parsed.data.windows ?? 0,
+    };
+}
 // Initialised once at module load — uses Application Default Credentials when
 // running on Cloud Run (or any GCP environment), and falls back to ADC from
 // the local environment during development.
@@ -30,28 +46,57 @@ async function readCounts(bucketName, objectPath) {
         return { ...DEFAULT_COUNTS };
     }
     const [contents] = await file.download();
-    const parsed = JSON.parse(contents.toString('utf8'));
+    return parseDownloadCounts(contents.toString('utf8'));
+}
+async function readCountsWithGeneration(bucketName, objectPath) {
+    const file = storage.bucket(bucketName).file(objectPath);
+    const [exists] = await file.exists();
+    if (!exists) {
+        return { counts: { ...DEFAULT_COUNTS }, generation: null, exists: false };
+    }
+    const [[metadata], [contents]] = await Promise.all([file.getMetadata(), file.download()]);
+    const counts = parseDownloadCounts(contents.toString('utf8'));
     return {
-        macos: typeof parsed.macos === 'number' ? parsed.macos : 0,
-        windows: typeof parsed.windows === 'number' ? parsed.windows : 0,
+        counts,
+        generation: metadata.generation ?? null,
+        exists: true,
     };
 }
-async function writeCounts(bucketName, objectPath, counts) {
-    const file = storage.bucket(bucketName).file(objectPath);
-    await file.save(JSON.stringify(counts), {
-        contentType: 'application/json',
-        resumable: false,
-    });
+function isGcsPreconditionError(err) {
+    const maybe = err;
+    return (maybe?.code === 412 ||
+        maybe?.message?.includes('conditionNotMet') === true ||
+        maybe?.message?.includes('Precondition Failed') === true);
 }
 async function incrementDownloadCount(platform) {
     const gcs = getGcsConfig();
     if (!gcs)
         return;
+    const file = storage.bucket(gcs.bucketName).file(gcs.objectPath);
+    const MAX_RETRIES = 6;
     try {
-        const counts = await readCounts(gcs.bucketName, gcs.objectPath);
-        counts[platform] += 1;
-        await writeCounts(gcs.bucketName, gcs.objectPath, counts);
-        logger_1.logger.info(`Download count incremented for ${platform}.`, { counts });
+        for (let attempt = 1; attempt <= MAX_RETRIES; attempt++) {
+            const { counts, generation, exists } = await readCountsWithGeneration(gcs.bucketName, gcs.objectPath);
+            counts[platform] += 1;
+            try {
+                await file.save(JSON.stringify(counts), {
+                    contentType: 'application/json',
+                    resumable: false,
+                    preconditionOpts: exists
+                        ? { ifGenerationMatch: Number(generation) }
+                        : { ifGenerationMatch: 0 },
+                });
+                logger_1.logger.info(`Download count incremented for ${platform}.`, { counts, attempt });
+                return;
+            }
+            catch (err) {
+                if (isGcsPreconditionError(err) && attempt < MAX_RETRIES) {
+                    continue;
+                }
+                throw err;
+            }
+        }
+        logger_1.logger.warn(`Download count increment exhausted retries for ${platform}.`);
     }
     catch (err) {
         logger_1.logger.error(`Failed to increment download count for ${platform}.`, { error: err });

package/backend-dist/index.js

@@ -8,6 +8,7 @@ const cors_1 = __importDefault(require("cors"));
 const path_1 = __importDefault(require("path"));
 const fs_1 = __importDefault(require("fs"));
 const zlib_1 = __importDefault(require("zlib"));
+const child_process_1 = require("child_process");
 const subscriptionRoutes_1 = require("./subscriptionRoutes");
 const featureRoutes_1 = require("./featureRoutes");
 const db_1 = require("./db");
@@ -23,6 +24,7 @@ require("./models/scheduledJob");
 const bucket_adapter_1 = require("./bucket-adapter");
 const app = (0, express_1.default)();
 const PORT = Number(config_1.config.port);
+const IS_CRON_CHILD = process.env.OMNIKEY_CRON_CHILD === '1';
 app.set('trust proxy', 1);
 app.use((0, cors_1.default)());
 app.use(express_1.default.json());
@@ -165,9 +167,38 @@ app.get('*', (_req, res) => {
     res.sendFile(path_1.default.join(process.cwd(), 'public', 'index.html'));
 });
 let server = null;
+let cronChildProcess = null;
+function startCronChildProcess() {
+    if (IS_CRON_CHILD || cronChildProcess)
+        return;
+    const childPort = PORT + 1;
+    const entry = process.argv[1] || __filename;
+    cronChildProcess = (0, child_process_1.fork)(entry, [], {
+        env: {
+            ...process.env,
+            OMNIKEY_CRON_CHILD: '1',
+            OMNIKEY_PORT: String(childPort),
+        },
+        execArgv: process.execArgv,
+        stdio: 'inherit',
+    });
+    logger_1.logger.info('Spawned cron child process.', {
+        pid: cronChildProcess.pid,
+        port: childPort,
+    });
+    cronChildProcess.on('exit', (code, signal) => {
+        logger_1.logger.warn('Cron child process exited.', { code, signal });
+        cronChildProcess = null;
+    });
+}
 async function start() {
     try {
         await (0, db_1.initDatabase)(logger_1.logger);
+        if (IS_CRON_CHILD) {
+            logger_1.logger.info('Starting cron child process mode.', { port: PORT });
+            (0, scheduledJobExecutor_1.startScheduledJobExecutor)();
+            return;
+        }
         server = app.listen(PORT, () => {
             logger_1.logger.info(`Enhancer API listening on http://localhost:${PORT}`, {
                 isSelfHosted: config_1.config.isSelfHosted,
@@ -181,7 +212,7 @@ async function start() {
             (0, agentServer_1.attachAgentWebSocketServer)(server);
         }
         if (config_1.config.isSelfHosted) {
-            (0, scheduledJobExecutor_1.startScheduledJobExecutor)();
+            startCronChildProcess();
         }
     }
     catch (err) {
@@ -192,6 +223,15 @@ async function start() {
 start();
 function gracefulShutdown(signal) {
     logger_1.logger.info(`Received ${signal}. Starting graceful shutdown...`);
+    if (cronChildProcess) {
+        cronChildProcess.kill('SIGTERM');
+        cronChildProcess = null;
+    }
+    if (IS_CRON_CHILD) {
+        logger_1.logger.info('Cron child process exiting.');
+        process.exit(0);
+        return;
+    }
     if (!server) {
         logger_1.logger.info('Server was not started or already closed. Exiting process.');
         process.exit(0);

package/package.json

@@ -4,7 +4,7 @@
     "access": "public",
     "registry": "https://registry.npmjs.org/"
   },
-  "version": "1.0.35",
+  "version": "1.0.36",
   "description": "CLI for onboarding users to Omnikey AI and configuring OPENAI_API_KEY. Use Yarn for install/build.",
   "engines": {
     "node": ">=14.0.0",