omnikey-cli 1.0.0 → 1.0.2
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +1 -1
- package/backend-dist/agentPrompts.js +64 -0
- package/backend-dist/agentServer.js +443 -0
- package/backend-dist/authMiddleware.js +85 -0
- package/backend-dist/compression.js +32 -0
- package/backend-dist/config.js +55 -0
- package/backend-dist/crypto.js +43 -0
- package/backend-dist/db.js +37 -0
- package/backend-dist/featureRoutes.js +296 -0
- package/backend-dist/index.js +123 -0
- package/backend-dist/logger.js +31 -0
- package/backend-dist/models/subscription.js +56 -0
- package/backend-dist/models/subscriptionTaskTemplate.js +58 -0
- package/backend-dist/models/subscriptionUsage.js +69 -0
- package/backend-dist/prompts.js +113 -0
- package/backend-dist/subscriptionRoutes.js +123 -0
- package/backend-dist/taskInstructionRoutes.js +157 -0
- package/backend-dist/types.js +10 -0
- package/dist/daemon.js +2 -2
- package/dist/index.js +4 -4
- package/dist/killDaemon.js +3 -3
- package/dist/onboard.js +1 -1
- package/package.json +7 -3
- package/src/daemon.ts +1 -1
- package/src/onboard.ts +3 -1
package/README.md
CHANGED
|
@@ -19,7 +19,7 @@ OmnikeyAI is a productivity tool for macOS that helps you quickly rewrite select
|
|
|
19
19
|
|
|
20
20
|
```sh
|
|
21
21
|
# Install CLI globally (from this directory)
|
|
22
|
-
npm install -g
|
|
22
|
+
npm install -g omnikey-cli
|
|
23
23
|
|
|
24
24
|
# Onboard interactively (will prompt for OpenAI key and self-hosting)
|
|
25
25
|
omnikey onboard
|
|
@@ -0,0 +1,64 @@
|
|
|
1
|
+
"use strict";
|
|
2
|
+
Object.defineProperty(exports, "__esModule", { value: true });
|
|
3
|
+
exports.AGENT_SYSTEM_PROMPT = void 0;
|
|
4
|
+
exports.AGENT_SYSTEM_PROMPT = `
|
|
5
|
+
You are an AI agent that can both reason about the user's situation and design shell scripts that the user will run on their own machine.
|
|
6
|
+
|
|
7
|
+
This agent is invoked when the user includes @omniAgent and there may also be stored custom task instructions for the current task.
|
|
8
|
+
Your job is to:
|
|
9
|
+
- Read and respect the stored task instructions (how to behave, what to focus on, output style) when they are provided.
|
|
10
|
+
- Carefully consider the current user input (what they typed when running @omniAgent).
|
|
11
|
+
- Decide whether additional machine-level information is needed, and if so, generate an appropriate shell script to gather it.
|
|
12
|
+
- Use the results of any previously run scripts plus the instructions and input to produce a complete, helpful final answer.
|
|
13
|
+
|
|
14
|
+
General guidelines:
|
|
15
|
+
- Only create commands that are safe and read-only, focusing on inspection, diagnostics, and information gathering.
|
|
16
|
+
- Do not generate any commands that install software, modify user data, or change system settings.
|
|
17
|
+
- Never ask the user to run commands with sudo or administrator/root privileges.
|
|
18
|
+
- Ensure that all commands provided are compatible with macOS and Linux; avoid any Windows-specific commands.
|
|
19
|
+
- Scripts must be self-contained and ready to run as-is, without the user needing to edit them.
|
|
20
|
+
|
|
21
|
+
The user will run the script and share the output with you.
|
|
22
|
+
|
|
23
|
+
<instruction_handling>
|
|
24
|
+
- Treat stored task instructions (if present) as authoritative for how to prioritize, what to examine, and how to format your answer, as long as they do not conflict with system rules or safety guidelines.
|
|
25
|
+
- Treat the current user input as the immediate goal or question you must solve, applying the stored instructions to that specific situation.
|
|
26
|
+
- If there is a conflict, follow: system rules first, then stored instructions, then ad-hoc guidance in the current input.
|
|
27
|
+
</instruction_handling>
|
|
28
|
+
|
|
29
|
+
<interaction_rules>
|
|
30
|
+
- When you need to execute ANY shell command, respond with a single <shell_script> block that contains the FULL script to run.
|
|
31
|
+
- Within that script, include all steps needed to carry out the current diagnostic or information-gathering task as completely as possible (for example, collect all relevant logs, inspect all relevant services, perform all necessary checks), rather than issuing minimal or placeholder commands.
|
|
32
|
+
- Prefer one comprehensive script over multiple small scripts; only wait for another round of output if you genuinely need the previous results to decide on the next actions.
|
|
33
|
+
- If further machine-level investigation is unnecessary, skip the shell script and respond directly with a <final_answer>.
|
|
34
|
+
- Every response MUST be exactly one of:
|
|
35
|
+
- A single <shell_script>...</shell_script> block, and nothing else; or
|
|
36
|
+
- A single <final_answer>...</final_answer> block, and nothing else.
|
|
37
|
+
- Never send plain text or explanation outside of these tags. If you are not emitting a <shell_script>, you MUST emit a <final_answer>.
|
|
38
|
+
- When you are completely finished and ready to present the result back to the user, respond with a single <final_answer> block.
|
|
39
|
+
- Do NOT include reasoning, commentary, or any other tags outside of <shell_script>...</shell_script> or <final_answer>...</final_answer>.
|
|
40
|
+
- Never wrap your entire response in other XML or JSON structures.
|
|
41
|
+
</interaction_rules>
|
|
42
|
+
|
|
43
|
+
<shell_script_block>
|
|
44
|
+
- Always emit exactly this structure when you want to run commands:
|
|
45
|
+
|
|
46
|
+
<shell_script>
|
|
47
|
+
#!/usr/bin/env bash
|
|
48
|
+
set -euo pipefail
|
|
49
|
+
# your commands here
|
|
50
|
+
</shell_script>
|
|
51
|
+
|
|
52
|
+
- Use a single, self-contained script per turn; do not send multiple <shell_script> blocks in one response.
|
|
53
|
+
- Inside the script, group related commands logically and add brief inline comments ONLY when they clarify non-obvious steps.
|
|
54
|
+
- Prefer safe, idempotent commands. Never ask for sudo.
|
|
55
|
+
</shell_script_block>
|
|
56
|
+
|
|
57
|
+
<final_answer_block>
|
|
58
|
+
- When you have gathered enough information and completed the requested work, respond once with:
|
|
59
|
+
<final_answer>
|
|
60
|
+
...user-facing result here (clear summary, key findings, concrete recommendations or next steps, formatted according to any stored instructions)...
|
|
61
|
+
</final_answer>
|
|
62
|
+
- Do not emit any text before or after the <final_answer> block; the entire response must be inside the <final_answer> tags.
|
|
63
|
+
</final_answer_block>
|
|
64
|
+
`;
|
|
@@ -0,0 +1,443 @@
|
|
|
1
|
+
"use strict";
|
|
2
|
+
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
|
|
3
|
+
if (k2 === undefined) k2 = k;
|
|
4
|
+
var desc = Object.getOwnPropertyDescriptor(m, k);
|
|
5
|
+
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
|
|
6
|
+
desc = { enumerable: true, get: function() { return m[k]; } };
|
|
7
|
+
}
|
|
8
|
+
Object.defineProperty(o, k2, desc);
|
|
9
|
+
}) : (function(o, m, k, k2) {
|
|
10
|
+
if (k2 === undefined) k2 = k;
|
|
11
|
+
o[k2] = m[k];
|
|
12
|
+
}));
|
|
13
|
+
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
|
|
14
|
+
Object.defineProperty(o, "default", { enumerable: true, value: v });
|
|
15
|
+
}) : function(o, v) {
|
|
16
|
+
o["default"] = v;
|
|
17
|
+
});
|
|
18
|
+
var __importStar = (this && this.__importStar) || (function () {
|
|
19
|
+
var ownKeys = function(o) {
|
|
20
|
+
ownKeys = Object.getOwnPropertyNames || function (o) {
|
|
21
|
+
var ar = [];
|
|
22
|
+
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
|
|
23
|
+
return ar;
|
|
24
|
+
};
|
|
25
|
+
return ownKeys(o);
|
|
26
|
+
};
|
|
27
|
+
return function (mod) {
|
|
28
|
+
if (mod && mod.__esModule) return mod;
|
|
29
|
+
var result = {};
|
|
30
|
+
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
|
|
31
|
+
__setModuleDefault(result, mod);
|
|
32
|
+
return result;
|
|
33
|
+
};
|
|
34
|
+
})();
|
|
35
|
+
var __importDefault = (this && this.__importDefault) || function (mod) {
|
|
36
|
+
return (mod && mod.__esModule) ? mod : { "default": mod };
|
|
37
|
+
};
|
|
38
|
+
Object.defineProperty(exports, "__esModule", { value: true });
|
|
39
|
+
exports.attachAgentWebSocketServer = attachAgentWebSocketServer;
|
|
40
|
+
const ws_1 = __importStar(require("ws"));
|
|
41
|
+
const jsonwebtoken_1 = __importDefault(require("jsonwebtoken"));
|
|
42
|
+
const openai_1 = __importDefault(require("openai"));
|
|
43
|
+
const cuid_1 = __importDefault(require("cuid"));
|
|
44
|
+
const config_1 = require("./config");
|
|
45
|
+
const logger_1 = require("./logger");
|
|
46
|
+
const subscription_1 = require("./models/subscription");
|
|
47
|
+
const subscriptionUsage_1 = require("./models/subscriptionUsage");
|
|
48
|
+
const agentPrompts_1 = require("./agentPrompts");
|
|
49
|
+
const featureRoutes_1 = require("./featureRoutes");
|
|
50
|
+
const authMiddleware_1 = require("./authMiddleware");
|
|
51
|
+
const openai = new openai_1.default({
|
|
52
|
+
apiKey: config_1.config.openaiApiKey,
|
|
53
|
+
});
|
|
54
|
+
const sessionMessages = new Map();
|
|
55
|
+
const MAX_TURNS = 10;
|
|
56
|
+
async function getOrCreateSession(sessionId, subscription, log) {
|
|
57
|
+
const existing = sessionMessages.get(sessionId);
|
|
58
|
+
if (existing) {
|
|
59
|
+
log.debug('Reusing existing agent session', {
|
|
60
|
+
sessionId,
|
|
61
|
+
subscriptionId: existing.subscription.id,
|
|
62
|
+
turns: existing.turns,
|
|
63
|
+
});
|
|
64
|
+
return existing;
|
|
65
|
+
}
|
|
66
|
+
// use these instructions as user instructions
|
|
67
|
+
const prompt = await (0, featureRoutes_1.getPromptForCommand)(log, 'task', subscription).catch((err) => {
|
|
68
|
+
log.error('Failed to get system prompt for new agent session', { error: err });
|
|
69
|
+
return '';
|
|
70
|
+
});
|
|
71
|
+
const entry = {
|
|
72
|
+
subscription,
|
|
73
|
+
history: [
|
|
74
|
+
{
|
|
75
|
+
role: 'system',
|
|
76
|
+
content: agentPrompts_1.AGENT_SYSTEM_PROMPT,
|
|
77
|
+
},
|
|
78
|
+
...(prompt
|
|
79
|
+
? [
|
|
80
|
+
{
|
|
81
|
+
role: 'assistant',
|
|
82
|
+
content: `<user_configured_instructions>
|
|
83
|
+
# User-Configured Task Instructions
|
|
84
|
+
${prompt}
|
|
85
|
+
</user_configured_instructions>`,
|
|
86
|
+
},
|
|
87
|
+
]
|
|
88
|
+
: []),
|
|
89
|
+
],
|
|
90
|
+
turns: 0,
|
|
91
|
+
};
|
|
92
|
+
sessionMessages.set(sessionId, entry);
|
|
93
|
+
log.info('Created new agent session', {
|
|
94
|
+
sessionId,
|
|
95
|
+
subscriptionId: subscription.id,
|
|
96
|
+
hasCustomPrompt: Boolean(prompt),
|
|
97
|
+
});
|
|
98
|
+
return entry;
|
|
99
|
+
}
|
|
100
|
+
async function authenticateFromAuthHeader(authHeader, log) {
|
|
101
|
+
if (config_1.config.isSelfHosted) {
|
|
102
|
+
log.info('Self-hosted mode: skipping JWT authentication for agent WebSocket connection.');
|
|
103
|
+
try {
|
|
104
|
+
const subscription = await (0, authMiddleware_1.selfHostedSubscription)();
|
|
105
|
+
log.info('Retrieved self-hosted subscription for agent WebSocket connection', {
|
|
106
|
+
subscriptionId: subscription.id,
|
|
107
|
+
});
|
|
108
|
+
return subscription;
|
|
109
|
+
}
|
|
110
|
+
catch (err) {
|
|
111
|
+
log.error('Failed to retrieve self-hosted subscription for agent WebSocket connection', {
|
|
112
|
+
error: err,
|
|
113
|
+
});
|
|
114
|
+
return null;
|
|
115
|
+
}
|
|
116
|
+
}
|
|
117
|
+
if (!config_1.config.jwtSecret) {
|
|
118
|
+
log.error('JWT secret is not configured. Cannot authenticate subscription from auth header.');
|
|
119
|
+
return null;
|
|
120
|
+
}
|
|
121
|
+
if (!authHeader) {
|
|
122
|
+
log.warn('Agent WebSocket connection missing authorization header');
|
|
123
|
+
return null;
|
|
124
|
+
}
|
|
125
|
+
const [scheme, token] = authHeader.split(' ');
|
|
126
|
+
if (scheme !== 'Bearer' || !token) {
|
|
127
|
+
log.warn('Agent WebSocket connection has malformed authorization header');
|
|
128
|
+
return null;
|
|
129
|
+
}
|
|
130
|
+
try {
|
|
131
|
+
const decoded = jsonwebtoken_1.default.verify(token, config_1.config.jwtSecret);
|
|
132
|
+
const subscription = await subscription_1.Subscription.findByPk(decoded.sid);
|
|
133
|
+
if (!subscription) {
|
|
134
|
+
log.warn('Agent WebSocket auth failed: subscription not found', {
|
|
135
|
+
sid: decoded.sid,
|
|
136
|
+
});
|
|
137
|
+
return null;
|
|
138
|
+
}
|
|
139
|
+
if (subscription.subscriptionStatus === 'expired') {
|
|
140
|
+
log.warn('Agent WebSocket auth failed: subscription expired', {
|
|
141
|
+
sid: decoded.sid,
|
|
142
|
+
});
|
|
143
|
+
return null;
|
|
144
|
+
}
|
|
145
|
+
const now = new Date();
|
|
146
|
+
if (subscription.licenseKeyExpiresAt && subscription.licenseKeyExpiresAt <= now) {
|
|
147
|
+
subscription.subscriptionStatus = 'expired';
|
|
148
|
+
await subscription.save();
|
|
149
|
+
log.info('Agent WebSocket auth: subscription key expired during connection', {
|
|
150
|
+
subscriptionId: subscription.id,
|
|
151
|
+
});
|
|
152
|
+
return null;
|
|
153
|
+
}
|
|
154
|
+
log.debug('Agent WebSocket auth succeeded', {
|
|
155
|
+
subscriptionId: subscription.id,
|
|
156
|
+
status: subscription.subscriptionStatus,
|
|
157
|
+
});
|
|
158
|
+
return subscription;
|
|
159
|
+
}
|
|
160
|
+
catch (err) {
|
|
161
|
+
log.warn('Agent WebSocket auth failed: invalid or expired JWT', { error: err });
|
|
162
|
+
return null;
|
|
163
|
+
}
|
|
164
|
+
}
|
|
165
|
+
async function runAgentTurn(sessionId, subscription, clientMessage, send, log) {
|
|
166
|
+
const session = await getOrCreateSession(sessionId, subscription, log);
|
|
167
|
+
// Count this call as one agent iteration.
|
|
168
|
+
session.turns += 1;
|
|
169
|
+
log.info('Starting agent turn', {
|
|
170
|
+
sessionId,
|
|
171
|
+
subscriptionId: subscription.id,
|
|
172
|
+
turn: session.turns,
|
|
173
|
+
});
|
|
174
|
+
// On the MAX_TURNS iteration, instruct the LLM to provide a final,
|
|
175
|
+
// consolidated answer based on the full conversation context.
|
|
176
|
+
if (session.turns === MAX_TURNS) {
|
|
177
|
+
session.history.push({
|
|
178
|
+
role: 'system',
|
|
179
|
+
content: 'Provide a single, final, concise answer based on the entire conversation so far. Wrap the answer in a <final_answer>...</final_answer> block and do not ask for further input or mention additional shell scripts to run. Do not include any <shell_script> block in this response.',
|
|
180
|
+
});
|
|
181
|
+
}
|
|
182
|
+
// Append the client message as user content, marking terminal
|
|
183
|
+
// output and errors in the text so the agent can reason about them.
|
|
184
|
+
let userContent = clientMessage.content || '';
|
|
185
|
+
const isTerminalOutput = Boolean(clientMessage.is_terminal_output);
|
|
186
|
+
const isErrorFlag = Boolean(clientMessage.is_error);
|
|
187
|
+
if (isTerminalOutput) {
|
|
188
|
+
userContent = `TERMINAL OUTPUT:\n${userContent}`;
|
|
189
|
+
}
|
|
190
|
+
if (isErrorFlag) {
|
|
191
|
+
userContent = `COMMAND ERROR:\n${userContent}`;
|
|
192
|
+
}
|
|
193
|
+
log.info('Agent turn received client message', {
|
|
194
|
+
sessionId,
|
|
195
|
+
isTerminalOutput,
|
|
196
|
+
isError: isErrorFlag,
|
|
197
|
+
rawContentLength: (clientMessage.content || '').length,
|
|
198
|
+
userContentLength: userContent.length,
|
|
199
|
+
});
|
|
200
|
+
session.history.push({
|
|
201
|
+
role: 'user',
|
|
202
|
+
content: userContent,
|
|
203
|
+
});
|
|
204
|
+
if (!config_1.config.openaiApiKey) {
|
|
205
|
+
log.warn('OPENAI_API_KEY is not set; returning error to client.');
|
|
206
|
+
const errorMessage = 'The server is missing its OpenAI API key. Please configure OPENAI_API_KEY on the backend and try again.';
|
|
207
|
+
send({
|
|
208
|
+
session_id: sessionId,
|
|
209
|
+
sender: 'agent',
|
|
210
|
+
content: `<final_answer>\n${errorMessage}\n</final_answer>`,
|
|
211
|
+
is_terminal_output: false,
|
|
212
|
+
is_error: true,
|
|
213
|
+
});
|
|
214
|
+
// Clear any cached session state so a subsequent attempt can
|
|
215
|
+
// start fresh once the environment is correctly configured.
|
|
216
|
+
sessionMessages.delete(sessionId);
|
|
217
|
+
return;
|
|
218
|
+
}
|
|
219
|
+
try {
|
|
220
|
+
log.debug('Calling OpenAI for agent turn', {
|
|
221
|
+
sessionId,
|
|
222
|
+
turn: session.turns,
|
|
223
|
+
historyLength: session.history.length,
|
|
224
|
+
});
|
|
225
|
+
const completion = await openai.chat.completions.create({
|
|
226
|
+
model: 'gpt-5.1',
|
|
227
|
+
// The OpenAI client accepts a superset of this simple
|
|
228
|
+
// message shape; we safely cast here to keep our local
|
|
229
|
+
// types minimal.
|
|
230
|
+
messages: session.history,
|
|
231
|
+
temperature: 0.2,
|
|
232
|
+
});
|
|
233
|
+
// Record token usage for this subscription and model, if usage
|
|
234
|
+
// data is available and we know which subscription made the call.
|
|
235
|
+
const usage = completion.usage;
|
|
236
|
+
if (usage && subscription.id) {
|
|
237
|
+
try {
|
|
238
|
+
await subscriptionUsage_1.SubscriptionUsage.create({
|
|
239
|
+
subscriptionId: subscription.id,
|
|
240
|
+
model: completion.model ?? 'gpt-5.1',
|
|
241
|
+
promptTokens: usage.prompt_tokens ?? 0,
|
|
242
|
+
completionTokens: usage.completion_tokens ?? 0,
|
|
243
|
+
totalTokens: usage.total_tokens ?? 0,
|
|
244
|
+
});
|
|
245
|
+
await subscription_1.Subscription.increment('totalTokensUsed', {
|
|
246
|
+
by: usage.total_tokens ?? 0,
|
|
247
|
+
where: { id: subscription.id },
|
|
248
|
+
});
|
|
249
|
+
}
|
|
250
|
+
catch (err) {
|
|
251
|
+
log.error('Failed to record subscription usage metrics for agent.', {
|
|
252
|
+
error: err,
|
|
253
|
+
subscriptionId: subscription.id,
|
|
254
|
+
});
|
|
255
|
+
}
|
|
256
|
+
}
|
|
257
|
+
const choice = completion.choices[0];
|
|
258
|
+
const content = (choice.message.content ?? '').toString().trim();
|
|
259
|
+
if (!content) {
|
|
260
|
+
log.warn('Agent LLM returned empty content; sending generic error to client.');
|
|
261
|
+
const errorMessage = 'The agent returned an empty response. Please try again.';
|
|
262
|
+
sendFinalAnswer(send, sessionId, errorMessage, true);
|
|
263
|
+
// Clear any cached session state so a subsequent attempt can
|
|
264
|
+
// start fresh without a polluted history.
|
|
265
|
+
sessionMessages.delete(sessionId);
|
|
266
|
+
return;
|
|
267
|
+
}
|
|
268
|
+
// Ensure that a proper <final_answer> block is produced for the
|
|
269
|
+
// desktop clients once we reach the final turn. If the model did
|
|
270
|
+
// not emit either a <shell_script> or <final_answer> tag on the
|
|
271
|
+
// MAX_TURNS turn, we treat this as the final natural-language answer
|
|
272
|
+
// and wrap it in <final_answer> tags so the client can stop
|
|
273
|
+
// waiting and paste the result.
|
|
274
|
+
const hasShellScriptTag = content.includes('<shell_script>');
|
|
275
|
+
const hasFinalAnswerTag = content.includes('<final_answer>');
|
|
276
|
+
log.info('Agent LLM raw response summary', {
|
|
277
|
+
sessionId,
|
|
278
|
+
turn: session.turns,
|
|
279
|
+
rawContentLength: content.length,
|
|
280
|
+
hasShellScriptTag,
|
|
281
|
+
hasFinalAnswerTag,
|
|
282
|
+
});
|
|
283
|
+
const normalizedContent = !hasShellScriptTag && !hasFinalAnswerTag && session.turns >= MAX_TURNS
|
|
284
|
+
? `<final_answer>\n${content}\n</final_answer>`
|
|
285
|
+
: content;
|
|
286
|
+
log.info('Agent LLM normalized response summary', {
|
|
287
|
+
sessionId,
|
|
288
|
+
turn: session.turns,
|
|
289
|
+
normalizedContentLength: normalizedContent.length,
|
|
290
|
+
});
|
|
291
|
+
// Record assistant message back into history for future turns.
|
|
292
|
+
session.history.push({
|
|
293
|
+
role: 'assistant',
|
|
294
|
+
content: normalizedContent,
|
|
295
|
+
});
|
|
296
|
+
send({
|
|
297
|
+
session_id: sessionId,
|
|
298
|
+
sender: 'agent',
|
|
299
|
+
content: normalizedContent,
|
|
300
|
+
is_terminal_output: false,
|
|
301
|
+
is_error: false,
|
|
302
|
+
});
|
|
303
|
+
// After the MAX_TURNS iteration or if a final answer tag is present, treat this as the final answer
|
|
304
|
+
// and clear the session from memory while marking it completed.
|
|
305
|
+
if (session.turns >= MAX_TURNS || hasFinalAnswerTag) {
|
|
306
|
+
log.info('Finalizing agent session after max turns or final answer tag', {
|
|
307
|
+
sessionId,
|
|
308
|
+
subscriptionId: subscription.id,
|
|
309
|
+
turns: session.turns,
|
|
310
|
+
hasFinalAnswerTag,
|
|
311
|
+
});
|
|
312
|
+
sessionMessages.delete(sessionId);
|
|
313
|
+
}
|
|
314
|
+
log.info('Completed agent turn', {
|
|
315
|
+
sessionId,
|
|
316
|
+
subscriptionId: subscription.id,
|
|
317
|
+
turn: session.turns,
|
|
318
|
+
responseLength: normalizedContent.length,
|
|
319
|
+
});
|
|
320
|
+
}
|
|
321
|
+
catch (err) {
|
|
322
|
+
log.error('Agent LLM call failed', { error: err });
|
|
323
|
+
const errorMessage = 'Agent failed to call language model. Please try again later.';
|
|
324
|
+
sendFinalAnswer(send, sessionId, errorMessage, true);
|
|
325
|
+
// Clear any cached session state so a subsequent attempt can
|
|
326
|
+
// start fresh without being polluted by a failed turn.
|
|
327
|
+
sessionMessages.delete(sessionId);
|
|
328
|
+
}
|
|
329
|
+
}
|
|
330
|
+
function sendFinalAnswer(send, sessionId, message, isError) {
|
|
331
|
+
send({
|
|
332
|
+
session_id: sessionId,
|
|
333
|
+
sender: 'agent',
|
|
334
|
+
content: `<final_answer>\n${message}\n</final_answer>`,
|
|
335
|
+
is_terminal_output: false,
|
|
336
|
+
is_error: isError,
|
|
337
|
+
});
|
|
338
|
+
}
|
|
339
|
+
function createLazyAuthContext(authHeader, log) {
|
|
340
|
+
let authenticatedSubscription = null;
|
|
341
|
+
let authFailed = false;
|
|
342
|
+
let authPromise = null;
|
|
343
|
+
const ensureAuthenticated = async () => {
|
|
344
|
+
if (authenticatedSubscription) {
|
|
345
|
+
return true;
|
|
346
|
+
}
|
|
347
|
+
if (authFailed) {
|
|
348
|
+
return false;
|
|
349
|
+
}
|
|
350
|
+
if (!authPromise) {
|
|
351
|
+
authPromise = (async () => {
|
|
352
|
+
try {
|
|
353
|
+
const sub = await authenticateFromAuthHeader(authHeader, log);
|
|
354
|
+
if (!sub) {
|
|
355
|
+
authFailed = true;
|
|
356
|
+
return;
|
|
357
|
+
}
|
|
358
|
+
authenticatedSubscription = sub;
|
|
359
|
+
log.info('Agent WebSocket authenticated', {
|
|
360
|
+
subscriptionId: authenticatedSubscription.id,
|
|
361
|
+
});
|
|
362
|
+
}
|
|
363
|
+
catch (err) {
|
|
364
|
+
authFailed = true;
|
|
365
|
+
log.error('Unexpected error during agent WebSocket auth', { error: err });
|
|
366
|
+
}
|
|
367
|
+
})();
|
|
368
|
+
}
|
|
369
|
+
await authPromise;
|
|
370
|
+
return Boolean(authenticatedSubscription);
|
|
371
|
+
};
|
|
372
|
+
const getSubscription = () => authenticatedSubscription;
|
|
373
|
+
return { ensureAuthenticated, getSubscription };
|
|
374
|
+
}
|
|
375
|
+
function attachAgentWebSocketServer(server) {
|
|
376
|
+
const wss = new ws_1.WebSocketServer({ server, path: '/ws/omni-agent' });
|
|
377
|
+
wss.on('connection', (ws, req) => {
|
|
378
|
+
const traceId = (0, cuid_1.default)();
|
|
379
|
+
const log = logger_1.logger.child({ traceId });
|
|
380
|
+
log.info('Agent WebSocket connection opened');
|
|
381
|
+
const authHeaderValue = req.headers['authorization'];
|
|
382
|
+
const authHeader = Array.isArray(authHeaderValue) ? authHeaderValue[0] : authHeaderValue;
|
|
383
|
+
const { ensureAuthenticated, getSubscription } = createLazyAuthContext(authHeader, log);
|
|
384
|
+
const send = (msg) => {
|
|
385
|
+
try {
|
|
386
|
+
ws.send(JSON.stringify(msg));
|
|
387
|
+
}
|
|
388
|
+
catch (err) {
|
|
389
|
+
log.error('Failed to write AgentMessage to WebSocket', { error: err });
|
|
390
|
+
}
|
|
391
|
+
};
|
|
392
|
+
ws.on('message', (data) => {
|
|
393
|
+
void (async () => {
|
|
394
|
+
const ok = await ensureAuthenticated();
|
|
395
|
+
const subscription = getSubscription();
|
|
396
|
+
if (!ok || !subscription) {
|
|
397
|
+
if (ws.readyState === ws_1.default.OPEN) {
|
|
398
|
+
log.warn('Closing Agent WebSocket due to failed authentication');
|
|
399
|
+
send({
|
|
400
|
+
session_id: '',
|
|
401
|
+
sender: 'agent',
|
|
402
|
+
content: 'Unauthorized: missing or invalid subscription. Please re-activate your key.',
|
|
403
|
+
is_terminal_output: false,
|
|
404
|
+
is_error: true,
|
|
405
|
+
});
|
|
406
|
+
ws.close();
|
|
407
|
+
}
|
|
408
|
+
return;
|
|
409
|
+
}
|
|
410
|
+
let message;
|
|
411
|
+
try {
|
|
412
|
+
const text = typeof data === 'string' ? data : data.toString('utf8');
|
|
413
|
+
log.info('Agent WebSocket received message from client', {
|
|
414
|
+
approximateLength: text.length,
|
|
415
|
+
});
|
|
416
|
+
message = JSON.parse(text);
|
|
417
|
+
}
|
|
418
|
+
catch (err) {
|
|
419
|
+
log.warn('Received invalid AgentMessage payload over WebSocket', { error: err });
|
|
420
|
+
return;
|
|
421
|
+
}
|
|
422
|
+
const sessionId = message.session_id || 'default';
|
|
423
|
+
log.debug('Received AgentMessage from client (WebSocket)', {
|
|
424
|
+
sessionId,
|
|
425
|
+
sender: message.sender,
|
|
426
|
+
isTerminalOutput: message.is_terminal_output,
|
|
427
|
+
isError: message.is_error,
|
|
428
|
+
});
|
|
429
|
+
void runAgentTurn(sessionId, subscription, message, send, log);
|
|
430
|
+
})();
|
|
431
|
+
});
|
|
432
|
+
ws.on('error', (err) => {
|
|
433
|
+
log.warn('Agent WebSocket error', { error: err });
|
|
434
|
+
});
|
|
435
|
+
ws.on('close', () => {
|
|
436
|
+
log.info('Agent WebSocket connection closed', {
|
|
437
|
+
hadAuthenticatedSubscription: Boolean(getSubscription()),
|
|
438
|
+
});
|
|
439
|
+
});
|
|
440
|
+
});
|
|
441
|
+
logger_1.logger.info('Agent WebSocket server attached at path /ws/omni-agent');
|
|
442
|
+
return wss;
|
|
443
|
+
}
|
|
@@ -0,0 +1,85 @@
|
|
|
1
|
+
"use strict";
|
|
2
|
+
var __importDefault = (this && this.__importDefault) || function (mod) {
|
|
3
|
+
return (mod && mod.__esModule) ? mod : { "default": mod };
|
|
4
|
+
};
|
|
5
|
+
Object.defineProperty(exports, "__esModule", { value: true });
|
|
6
|
+
exports.selfHostedSubscription = selfHostedSubscription;
|
|
7
|
+
exports.authMiddleware = authMiddleware;
|
|
8
|
+
const crypto_1 = require("crypto");
|
|
9
|
+
const jsonwebtoken_1 = __importDefault(require("jsonwebtoken"));
|
|
10
|
+
const logger_1 = require("./logger");
|
|
11
|
+
const config_1 = require("./config");
|
|
12
|
+
const subscription_1 = require("./models/subscription");
|
|
13
|
+
async function selfHostedSubscription() {
|
|
14
|
+
try {
|
|
15
|
+
let subscription = await subscription_1.Subscription.findOne({ where: { isSelfHosted: true } });
|
|
16
|
+
if (!subscription) {
|
|
17
|
+
subscription = await subscription_1.Subscription.create({
|
|
18
|
+
email: 'local-user@omnikey.ai',
|
|
19
|
+
licenseKey: 'self-hosted',
|
|
20
|
+
subscriptionStatus: 'active',
|
|
21
|
+
isSelfHosted: true,
|
|
22
|
+
});
|
|
23
|
+
logger_1.logger.info('Created self-hosted subscription record in database.');
|
|
24
|
+
}
|
|
25
|
+
return subscription;
|
|
26
|
+
}
|
|
27
|
+
catch (err) {
|
|
28
|
+
logger_1.logger.error('Error ensuring self-hosted subscription record exists.', { error: err });
|
|
29
|
+
throw err;
|
|
30
|
+
}
|
|
31
|
+
}
|
|
32
|
+
async function authMiddleware(req, res, next) {
|
|
33
|
+
const authHeader = req.headers.authorization;
|
|
34
|
+
logger_1.logger.defaultMeta = { traceId: (0, crypto_1.randomUUID)() };
|
|
35
|
+
if (config_1.config.isSelfHosted || !config_1.config.jwtSecret) {
|
|
36
|
+
logger_1.logger.info('Self-hosted mode: skipping auth middleware.');
|
|
37
|
+
if (config_1.config.isSelfHosted) {
|
|
38
|
+
res.locals.subscription = await selfHostedSubscription();
|
|
39
|
+
res.locals.logger = logger_1.logger;
|
|
40
|
+
}
|
|
41
|
+
return next();
|
|
42
|
+
}
|
|
43
|
+
if (!authHeader) {
|
|
44
|
+
logger_1.logger.warn('Missing Authorization header on feature route.');
|
|
45
|
+
return res.status(401).json({ error: 'Missing bearer token.' });
|
|
46
|
+
}
|
|
47
|
+
const [scheme, token] = authHeader.split(' ');
|
|
48
|
+
if (scheme !== 'Bearer' || !token) {
|
|
49
|
+
logger_1.logger.warn('Malformed Authorization header on feature route.');
|
|
50
|
+
return res.status(401).json({ error: 'Invalid authorization header.' });
|
|
51
|
+
}
|
|
52
|
+
try {
|
|
53
|
+
const decoded = jsonwebtoken_1.default.verify(token, config_1.config.jwtSecret);
|
|
54
|
+
const subscription = await subscription_1.Subscription.findByPk(decoded.sid);
|
|
55
|
+
if (!subscription) {
|
|
56
|
+
logger_1.logger.warn('Subscription not found for JWT.', { sid: decoded.sid });
|
|
57
|
+
return res.status(403).json({ error: 'Invalid or expired token.' });
|
|
58
|
+
}
|
|
59
|
+
if (subscription.subscriptionStatus == 'expired') {
|
|
60
|
+
logger_1.logger.warn('Inactive subscription for JWT.', {
|
|
61
|
+
sid: decoded.sid,
|
|
62
|
+
status: subscription.subscriptionStatus,
|
|
63
|
+
});
|
|
64
|
+
return res.status(403).json({ error: 'Subscription is not active.' });
|
|
65
|
+
}
|
|
66
|
+
const now = new Date();
|
|
67
|
+
if (subscription.licenseKeyExpiresAt && subscription.licenseKeyExpiresAt <= now) {
|
|
68
|
+
subscription.subscriptionStatus = 'expired';
|
|
69
|
+
await subscription.save();
|
|
70
|
+
logger_1.logger.info('Subscription key has expired during activation.', {
|
|
71
|
+
subscriptionId: subscription.id,
|
|
72
|
+
});
|
|
73
|
+
return res
|
|
74
|
+
.status(403)
|
|
75
|
+
.json({ error: 'Subscription has expired.', subscriptionStatus: 'expired' });
|
|
76
|
+
}
|
|
77
|
+
res.locals.logger = logger_1.logger;
|
|
78
|
+
res.locals.subscription = subscription;
|
|
79
|
+
next();
|
|
80
|
+
}
|
|
81
|
+
catch (err) {
|
|
82
|
+
logger_1.logger.warn('Invalid or expired JWT on feature route.', { error: err });
|
|
83
|
+
return res.status(403).json({ error: 'Invalid or expired token.' });
|
|
84
|
+
}
|
|
85
|
+
}
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
"use strict";
|
|
2
|
+
var __importDefault = (this && this.__importDefault) || function (mod) {
|
|
3
|
+
return (mod && mod.__esModule) ? mod : { "default": mod };
|
|
4
|
+
};
|
|
5
|
+
Object.defineProperty(exports, "__esModule", { value: true });
|
|
6
|
+
exports.compressString = compressString;
|
|
7
|
+
exports.decompressString = decompressString;
|
|
8
|
+
const zlib_1 = __importDefault(require("zlib"));
|
|
9
|
+
const COMPRESSED_PREFIX = 'gz1:';
|
|
10
|
+
function compressString(value) {
|
|
11
|
+
const buffer = Buffer.from(value, 'utf8');
|
|
12
|
+
const compressed = zlib_1.default.gzipSync(buffer);
|
|
13
|
+
return COMPRESSED_PREFIX + compressed.toString('base64');
|
|
14
|
+
}
|
|
15
|
+
function decompressString(value) {
|
|
16
|
+
if (value == null)
|
|
17
|
+
return null;
|
|
18
|
+
if (!value.startsWith(COMPRESSED_PREFIX)) {
|
|
19
|
+
// Backwards compatibility: treat as plain text.
|
|
20
|
+
return value;
|
|
21
|
+
}
|
|
22
|
+
try {
|
|
23
|
+
const b64 = value.slice(COMPRESSED_PREFIX.length);
|
|
24
|
+
const compressed = Buffer.from(b64, 'base64');
|
|
25
|
+
const decompressed = zlib_1.default.gunzipSync(compressed);
|
|
26
|
+
return decompressed.toString('utf8');
|
|
27
|
+
}
|
|
28
|
+
catch {
|
|
29
|
+
// If decompression fails, treat as missing instructions.
|
|
30
|
+
return null;
|
|
31
|
+
}
|
|
32
|
+
}
|