@xalia/agent 0.6.1 → 0.6.3
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/dist/agent/src/agent/agent.js +109 -57
- package/dist/agent/src/agent/agentUtils.js +24 -26
- package/dist/agent/src/agent/compressingContextManager.js +3 -2
- package/dist/agent/src/agent/dummyLLM.js +1 -3
- package/dist/agent/src/agent/imageGenLLM.js +67 -0
- package/dist/agent/src/agent/imageGenerator.js +43 -0
- package/dist/agent/src/agent/llm.js +27 -0
- package/dist/agent/src/agent/mcpServerManager.js +18 -6
- package/dist/agent/src/agent/nullAgentEventHandler.js +6 -0
- package/dist/agent/src/agent/openAILLM.js +3 -3
- package/dist/agent/src/agent/openAILLMStreaming.js +41 -6
- package/dist/agent/src/chat/client/chatClient.js +154 -235
- package/dist/agent/src/chat/client/constants.js +1 -2
- package/dist/agent/src/chat/client/sessionClient.js +47 -15
- package/dist/agent/src/chat/client/sessionFiles.js +102 -0
- package/dist/agent/src/chat/data/apiKeyManager.js +38 -7
- package/dist/agent/src/chat/data/database.js +83 -70
- package/dist/agent/src/chat/data/dbSessionFileModels.js +49 -0
- package/dist/agent/src/chat/data/dbSessionFiles.js +76 -0
- package/dist/agent/src/chat/data/dbSessionMessages.js +57 -0
- package/dist/agent/src/chat/data/mimeTypes.js +44 -0
- package/dist/agent/src/chat/protocol/messages.js +21 -1
- package/dist/agent/src/chat/server/chatContextManager.js +19 -16
- package/dist/agent/src/chat/server/connectionManager.js +14 -36
- package/dist/agent/src/chat/server/connectionManager.test.js +3 -16
- package/dist/agent/src/chat/server/conversation.js +73 -44
- package/dist/agent/src/chat/server/imageGeneratorTools.js +111 -0
- package/dist/agent/src/chat/server/openSession.js +398 -233
- package/dist/agent/src/chat/server/openSessionMessageSender.js +2 -0
- package/dist/agent/src/chat/server/server.js +5 -8
- package/dist/agent/src/chat/server/sessionFileManager.js +171 -38
- package/dist/agent/src/chat/server/sessionRegistry.js +214 -42
- package/dist/agent/src/chat/server/test-utils/mockFactories.js +12 -11
- package/dist/agent/src/chat/server/tools.js +27 -6
- package/dist/agent/src/chat/utils/approvalManager.js +82 -64
- package/dist/agent/src/chat/utils/multiAsyncQueue.js +9 -1
- package/dist/agent/src/chat/{client/responseHandler.js → utils/responseAwaiter.js} +41 -18
- package/dist/agent/src/test/agent.test.js +104 -63
- package/dist/agent/src/test/approvalManager.test.js +79 -35
- package/dist/agent/src/test/chatContextManager.test.js +16 -17
- package/dist/agent/src/test/clientServerConnection.test.js +2 -2
- package/dist/agent/src/test/db.test.js +33 -70
- package/dist/agent/src/test/dbSessionFiles.test.js +179 -0
- package/dist/agent/src/test/dbSessionMessages.test.js +67 -0
- package/dist/agent/src/test/dbTestTools.js +6 -5
- package/dist/agent/src/test/imageLoad.test.js +1 -1
- package/dist/agent/src/test/mcpServerManager.test.js +1 -1
- package/dist/agent/src/test/multiAsyncQueue.test.js +50 -0
- package/dist/agent/src/test/responseAwaiter.test.js +74 -0
- package/dist/agent/src/test/testTools.js +12 -0
- package/dist/agent/src/tool/agentChat.js +25 -6
- package/dist/agent/src/tool/agentMain.js +1 -1
- package/dist/agent/src/tool/chatMain.js +115 -6
- package/dist/agent/src/tool/commandPrompt.js +7 -3
- package/dist/agent/src/tool/files.js +23 -15
- package/dist/agent/src/tool/options.js +2 -2
- package/package.json +1 -1
- package/scripts/setup_chat +2 -2
- package/scripts/test_chat +95 -36
- package/src/agent/agent.ts +152 -41
- package/src/agent/agentUtils.ts +34 -41
- package/src/agent/compressingContextManager.ts +5 -4
- package/src/agent/context.ts +1 -1
- package/src/agent/dummyLLM.ts +1 -3
- package/src/agent/iAgentEventHandler.ts +15 -2
- package/src/agent/imageGenLLM.ts +99 -0
- package/src/agent/imageGenerator.ts +60 -0
- package/src/agent/llm.ts +128 -4
- package/src/agent/mcpServerManager.ts +26 -7
- package/src/agent/nullAgentEventHandler.ts +6 -0
- package/src/agent/openAILLM.ts +3 -8
- package/src/agent/openAILLMStreaming.ts +60 -14
- package/src/chat/client/chatClient.ts +262 -286
- package/src/chat/client/constants.ts +0 -2
- package/src/chat/client/sessionClient.ts +82 -20
- package/src/chat/client/sessionFiles.ts +142 -0
- package/src/chat/data/apiKeyManager.ts +55 -7
- package/src/chat/data/dataModels.ts +17 -7
- package/src/chat/data/database.ts +107 -92
- package/src/chat/data/dbSessionFileModels.ts +91 -0
- package/src/chat/data/dbSessionFiles.ts +99 -0
- package/src/chat/data/dbSessionMessages.ts +68 -0
- package/src/chat/data/mimeTypes.ts +58 -0
- package/src/chat/protocol/messages.ts +136 -25
- package/src/chat/server/chatContextManager.ts +42 -24
- package/src/chat/server/connectionManager.test.ts +2 -22
- package/src/chat/server/connectionManager.ts +18 -53
- package/src/chat/server/conversation.ts +106 -59
- package/src/chat/server/imageGeneratorTools.ts +138 -0
- package/src/chat/server/openSession.ts +606 -325
- package/src/chat/server/openSessionMessageSender.ts +4 -0
- package/src/chat/server/server.ts +5 -11
- package/src/chat/server/sessionFileManager.ts +223 -63
- package/src/chat/server/sessionRegistry.ts +317 -52
- package/src/chat/server/test-utils/mockFactories.ts +13 -13
- package/src/chat/server/tools.ts +43 -8
- package/src/chat/utils/agentSessionMap.ts +2 -2
- package/src/chat/utils/approvalManager.ts +153 -81
- package/src/chat/utils/multiAsyncQueue.ts +11 -1
- package/src/chat/{client/responseHandler.ts → utils/responseAwaiter.ts} +73 -23
- package/src/test/agent.test.ts +152 -75
- package/src/test/approvalManager.test.ts +108 -40
- package/src/test/chatContextManager.test.ts +26 -22
- package/src/test/clientServerConnection.test.ts +3 -3
- package/src/test/compressingContextManager.test.ts +1 -1
- package/src/test/context.test.ts +2 -1
- package/src/test/conversation.test.ts +1 -1
- package/src/test/db.test.ts +41 -83
- package/src/test/dbSessionFiles.test.ts +258 -0
- package/src/test/dbSessionMessages.test.ts +85 -0
- package/src/test/dbTestTools.ts +9 -5
- package/src/test/imageLoad.test.ts +2 -2
- package/src/test/mcpServerManager.test.ts +3 -1
- package/src/test/multiAsyncQueue.test.ts +58 -0
- package/src/test/responseAwaiter.test.ts +103 -0
- package/src/test/testTools.ts +15 -1
- package/src/tool/agentChat.ts +36 -8
- package/src/tool/agentMain.ts +7 -7
- package/src/tool/chatMain.ts +128 -7
- package/src/tool/commandPrompt.ts +10 -5
- package/src/tool/files.ts +30 -13
- package/src/tool/options.ts +1 -1
- package/test_data/dummyllm_script_image_gen.json +19 -0
- package/test_data/dummyllm_script_invoke_image_gen_tool.json +30 -0
- package/test_data/image_gen_test_profile.json +5 -0
- package/dist/agent/src/test/responseHandler.test.js +0 -61
- package/src/test/responseHandler.test.ts +0 -78
|
@@ -0,0 +1,57 @@
|
|
|
1
|
+
"use strict";
|
|
2
|
+
Object.defineProperty(exports, "__esModule", { value: true });
|
|
3
|
+
exports.DbSessionMessages = void 0;
|
|
4
|
+
const database_1 = require("./database");
|
|
5
|
+
class DbSessionMessages extends database_1.DbClientBase {
|
|
6
|
+
async clearConversation(session_uuid) {
|
|
7
|
+
const { error } = await this.client
|
|
8
|
+
.from("session_messages")
|
|
9
|
+
.delete()
|
|
10
|
+
.eq("session_uuid", session_uuid);
|
|
11
|
+
if (error) {
|
|
12
|
+
throw error;
|
|
13
|
+
}
|
|
14
|
+
}
|
|
15
|
+
async getConversation(session_uuid, numEntries, beforeIndex) {
|
|
16
|
+
// Query all message for the given session, ordered high-to-low by
|
|
17
|
+
// message_idx, limited to `numEntries`. If `beforeIndex` is given, it
|
|
18
|
+
// means we get messages with `message_idx < beforeIndex`
|
|
19
|
+
let query = this.client
|
|
20
|
+
.from("session_messages")
|
|
21
|
+
.select("message_idx,sender_uuid,is_for_llm,content")
|
|
22
|
+
.eq("session_uuid", session_uuid);
|
|
23
|
+
if (beforeIndex) {
|
|
24
|
+
query = query.lt("message_idx", beforeIndex);
|
|
25
|
+
}
|
|
26
|
+
query = query.order("message_idx", { ascending: false }).limit(numEntries);
|
|
27
|
+
const { data, error } = await query;
|
|
28
|
+
if (error) {
|
|
29
|
+
throw error;
|
|
30
|
+
}
|
|
31
|
+
// To get the newest N messages, we've orded by index largest to smallest
|
|
32
|
+
// (newest first), but caller wants the message first to last, hence
|
|
33
|
+
// reverse the array.
|
|
34
|
+
return data
|
|
35
|
+
.map(({ sender_uuid, ...rest }) => {
|
|
36
|
+
return typeof sender_uuid === "string"
|
|
37
|
+
? {
|
|
38
|
+
sender_uuid,
|
|
39
|
+
...rest,
|
|
40
|
+
}
|
|
41
|
+
: { ...rest };
|
|
42
|
+
})
|
|
43
|
+
.reverse();
|
|
44
|
+
}
|
|
45
|
+
async append(session_uuid, messages) {
|
|
46
|
+
const payload = messages.map((m) => {
|
|
47
|
+
return { ...m, session_uuid };
|
|
48
|
+
});
|
|
49
|
+
const { error } = await this.client
|
|
50
|
+
.from("session_messages")
|
|
51
|
+
.insert(payload);
|
|
52
|
+
if (error) {
|
|
53
|
+
throw error;
|
|
54
|
+
}
|
|
55
|
+
}
|
|
56
|
+
}
|
|
57
|
+
exports.DbSessionMessages = DbSessionMessages;
|
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
"use strict";
|
|
2
|
+
Object.defineProperty(exports, "__esModule", { value: true });
|
|
3
|
+
exports.EXTENSION_TO_IMAGE_MIME_TYPE = exports.IMAGE_MIME_TYPES = void 0;
|
|
4
|
+
exports.isImageMimeType = isImageMimeType;
|
|
5
|
+
exports.getMimeTypeFromDataUrl = getMimeTypeFromDataUrl;
|
|
6
|
+
exports.createDataUrlFromBuffer = createDataUrlFromBuffer;
|
|
7
|
+
exports.createDataUrlFromText = createDataUrlFromText;
|
|
8
|
+
const assert_1 = require("assert");
|
|
9
|
+
exports.IMAGE_MIME_TYPES = [
|
|
10
|
+
"image/png",
|
|
11
|
+
"image/jpeg",
|
|
12
|
+
"image/gif",
|
|
13
|
+
"image/webp",
|
|
14
|
+
];
|
|
15
|
+
function isImageMimeType(mime_type) {
|
|
16
|
+
return (typeof mime_type === "string" &&
|
|
17
|
+
exports.IMAGE_MIME_TYPES.includes(mime_type));
|
|
18
|
+
}
|
|
19
|
+
exports.EXTENSION_TO_IMAGE_MIME_TYPE = {
|
|
20
|
+
jpg: "image/jpeg",
|
|
21
|
+
jpeg: "image/jpeg",
|
|
22
|
+
png: "image/png",
|
|
23
|
+
gif: "image/gif",
|
|
24
|
+
webp: "image/webp",
|
|
25
|
+
};
|
|
26
|
+
function getMimeTypeFromDataUrl(data_url) {
|
|
27
|
+
// data:image/png;base64,AAAAA...
|
|
28
|
+
(0, assert_1.strict)(data_url.startsWith("data:"));
|
|
29
|
+
const endOfMimeType = data_url.indexOf(",");
|
|
30
|
+
(0, assert_1.strict)(endOfMimeType > 5);
|
|
31
|
+
let mimeType = data_url.slice(5, endOfMimeType);
|
|
32
|
+
const paramIdx = mimeType.indexOf(";");
|
|
33
|
+
if (paramIdx >= 0) {
|
|
34
|
+
mimeType = mimeType.slice(0, paramIdx);
|
|
35
|
+
}
|
|
36
|
+
return mimeType;
|
|
37
|
+
}
|
|
38
|
+
function createDataUrlFromBuffer(data, mime_type) {
|
|
39
|
+
const imgB64 = data.toString("base64");
|
|
40
|
+
return `data:${mime_type};base64,${imgB64}`;
|
|
41
|
+
}
|
|
42
|
+
function createDataUrlFromText(data, mime_type) {
|
|
43
|
+
return `data:${mime_type},${data}`;
|
|
44
|
+
}
|
|
@@ -2,6 +2,7 @@
|
|
|
2
2
|
Object.defineProperty(exports, "__esModule", { value: true });
|
|
3
3
|
exports.isClientControlMessage = isClientControlMessage;
|
|
4
4
|
exports.isServerControlMessage = isServerControlMessage;
|
|
5
|
+
exports.isServerSessionFileMessage = isServerSessionFileMessage;
|
|
5
6
|
exports.isServerSessionScopedMessage = isServerSessionScopedMessage;
|
|
6
7
|
exports.decodeAssistantMessageParam = decodeAssistantMessageParam;
|
|
7
8
|
function isClientControlMessage(message) {
|
|
@@ -35,6 +36,7 @@ function isServerControlMessage(message) {
|
|
|
35
36
|
case "control_session_left":
|
|
36
37
|
case "control_session_deleted":
|
|
37
38
|
case "control_team_created":
|
|
39
|
+
case "control_team_members_updated":
|
|
38
40
|
case "control_error":
|
|
39
41
|
return true;
|
|
40
42
|
default: {
|
|
@@ -43,6 +45,19 @@ function isServerControlMessage(message) {
|
|
|
43
45
|
}
|
|
44
46
|
}
|
|
45
47
|
}
|
|
48
|
+
function isServerSessionFileMessage(message) {
|
|
49
|
+
const msg = message;
|
|
50
|
+
switch (msg.type) {
|
|
51
|
+
case "session_file_changed":
|
|
52
|
+
case "session_file_deleted":
|
|
53
|
+
case "session_file_content":
|
|
54
|
+
return true;
|
|
55
|
+
default: {
|
|
56
|
+
const _ = msg;
|
|
57
|
+
return false;
|
|
58
|
+
}
|
|
59
|
+
}
|
|
60
|
+
}
|
|
46
61
|
/**
|
|
47
62
|
* Type guard to check if a ServerToClient message is a session-scoped message
|
|
48
63
|
*/
|
|
@@ -54,9 +69,9 @@ function isServerSessionScopedMessage(message) {
|
|
|
54
69
|
case "user_msg":
|
|
55
70
|
case "agent_msg":
|
|
56
71
|
case "agent_msg_chunk":
|
|
72
|
+
case "agent_reasoning_chunk":
|
|
57
73
|
case "user_joined":
|
|
58
74
|
case "user_left":
|
|
59
|
-
case "session_update":
|
|
60
75
|
case "tool_auto_approval_set":
|
|
61
76
|
case "tool_call":
|
|
62
77
|
case "tool_call_approval_result":
|
|
@@ -64,17 +79,22 @@ function isServerSessionScopedMessage(message) {
|
|
|
64
79
|
case "authentication_started":
|
|
65
80
|
case "authentication_finished":
|
|
66
81
|
case "user_typing":
|
|
82
|
+
case "session_file_changed":
|
|
83
|
+
case "session_file_deleted":
|
|
84
|
+
case "session_file_content":
|
|
67
85
|
case "mcp_server_added":
|
|
68
86
|
case "mcp_server_removed":
|
|
69
87
|
case "mcp_server_tool_enabled":
|
|
70
88
|
case "mcp_server_tool_disabled":
|
|
71
89
|
case "system_prompt_updated":
|
|
72
90
|
case "model_updated":
|
|
91
|
+
case "agent_paused":
|
|
73
92
|
case "user_added":
|
|
74
93
|
case "user_removed":
|
|
75
94
|
case "authenticate":
|
|
76
95
|
case "approve_tool_call":
|
|
77
96
|
case "render_html":
|
|
97
|
+
case "session_shared":
|
|
78
98
|
return true;
|
|
79
99
|
default: {
|
|
80
100
|
const _ = msg;
|
|
@@ -31,7 +31,7 @@ class ChatContextManager {
|
|
|
31
31
|
logger.debug(`[ChatContextManager]: llm messages: ${JSON.stringify(llmMessages)}`);
|
|
32
32
|
// Insert a system message placeholder into the context
|
|
33
33
|
this.sessionUUID = sessionUUID;
|
|
34
|
-
this.conversationMessages = (0, conversation_1.
|
|
34
|
+
this.conversationMessages = (0, conversation_1.sessionMessagesToConversationMessages)(sessionMessages, defaultUserName, sessionUUID);
|
|
35
35
|
this.pendingUserMessages = [];
|
|
36
36
|
this.llmContext = new compressingContextManager_1.CompressingContextManager(systemPrompt, llmMessages, compressionAgentUrl, compressionAgentModel, compressionAgentApiKey);
|
|
37
37
|
this.nextMessageIdx = nextMessageIdx;
|
|
@@ -41,7 +41,7 @@ class ChatContextManager {
|
|
|
41
41
|
this.pendingCompression = false;
|
|
42
42
|
this.checkpointWriter = checkpointWriter;
|
|
43
43
|
this.fileManager = fileManager;
|
|
44
|
-
fileManager.
|
|
44
|
+
fileManager.addEventHandler(this);
|
|
45
45
|
this.fileManagerDescriptionsDirty = true;
|
|
46
46
|
}
|
|
47
47
|
// IContextManager.addMessages
|
|
@@ -83,7 +83,10 @@ class ChatContextManager {
|
|
|
83
83
|
this.fileManagerDescriptionsDirty = true;
|
|
84
84
|
}
|
|
85
85
|
// ISessionFileManagerEventHandler.onFileChange
|
|
86
|
-
|
|
86
|
+
onFileChanged(_entry) {
|
|
87
|
+
this.fileManagerDescriptionsDirty = true;
|
|
88
|
+
}
|
|
89
|
+
onFileDeleted(_name) {
|
|
87
90
|
this.fileManagerDescriptionsDirty = true;
|
|
88
91
|
}
|
|
89
92
|
setWorkspace(userMessage) {
|
|
@@ -96,16 +99,21 @@ class ChatContextManager {
|
|
|
96
99
|
getConversationMessages() {
|
|
97
100
|
return this.conversationMessages.concat(this.pendingUserMessages);
|
|
98
101
|
}
|
|
99
|
-
processUserMessage(msg,
|
|
102
|
+
processUserMessage(msg, from_uuid, from_nickname) {
|
|
100
103
|
// TODO: maintain a queue internally instead of relying on the caller to
|
|
101
104
|
// pass in our generated messages back into `startAgentResponse`.
|
|
105
|
+
// Filter out null messages immediately.
|
|
106
|
+
if (!msg.imageB64 && !msg.message) {
|
|
107
|
+
return undefined;
|
|
108
|
+
}
|
|
102
109
|
const message_idx = this.getNextMessageIdx();
|
|
103
110
|
const userMessage = {
|
|
104
111
|
type: "user_msg",
|
|
105
112
|
session_id: this.sessionUUID,
|
|
106
113
|
message_idx,
|
|
107
114
|
message: msg.message,
|
|
108
|
-
user_uuid:
|
|
115
|
+
user_uuid: from_uuid,
|
|
116
|
+
user_nickname: from_nickname,
|
|
109
117
|
};
|
|
110
118
|
if (msg.imageB64) {
|
|
111
119
|
userMessage.imageB64 = msg.imageB64;
|
|
@@ -142,7 +150,7 @@ class ChatContextManager {
|
|
|
142
150
|
// Compute the new llm messages
|
|
143
151
|
const llmUserMessages = [];
|
|
144
152
|
for (const msg of pendingUserMessages) {
|
|
145
|
-
const userMsg = (0, agent_1.createUserMessage)(msg.message, msg.imageB64, msg.
|
|
153
|
+
const userMsg = (0, agent_1.createUserMessage)(msg.message, msg.imageB64, msg.user_nickname);
|
|
146
154
|
if (userMsg) {
|
|
147
155
|
llmUserMessages.push(userMsg);
|
|
148
156
|
}
|
|
@@ -159,11 +167,11 @@ class ChatContextManager {
|
|
|
159
167
|
};
|
|
160
168
|
}
|
|
161
169
|
endAgentResponse() {
|
|
162
|
-
(0, assert_1.strict)(typeof this.startingLLMContextLength !== "undefined");
|
|
163
|
-
(0, assert_1.strict)(typeof this.pendingMessages !== "undefined");
|
|
164
|
-
(0, assert_1.strict)(typeof this.curAgentMsgIdx !== "undefined");
|
|
170
|
+
(0, assert_1.strict)(typeof this.startingLLMContextLength !== "undefined", "agent response not started (startingLLMContextLength)");
|
|
171
|
+
(0, assert_1.strict)(typeof this.pendingMessages !== "undefined", "agent response not started (pendingMessages)");
|
|
172
|
+
(0, assert_1.strict)(typeof this.curAgentMsgIdx !== "undefined", "agent response not started (curAgentMsgIdx)");
|
|
165
173
|
const numPending = this.pendingMessages.length;
|
|
166
|
-
(0, assert_1.strict)(numPending >
|
|
174
|
+
(0, assert_1.strict)(numPending > 0, "no pending"); // at least 1 user message
|
|
167
175
|
// Compute DB messages
|
|
168
176
|
const newSessionMessages = (0, conversation_1.chatMessagesToSessionMessages)(this.pendingMessages);
|
|
169
177
|
const newLLMMessages = this.llmContext.getPending();
|
|
@@ -193,7 +201,7 @@ class ChatContextManager {
|
|
|
193
201
|
}
|
|
194
202
|
if (JSON.stringify(sMsg.content) !== JSON.stringify(lMsg)) {
|
|
195
203
|
messageListError(`newSessionMessages[${String(i)}].content !== ` +
|
|
196
|
-
`newLLMMessages[${String(i)}]
|
|
204
|
+
`newLLMMessages[${String(i)}]`);
|
|
197
205
|
}
|
|
198
206
|
if (sMsg.message_idx !== pMsg.message_idx) {
|
|
199
207
|
messageListError(`newSessionMessages[${String(i)}].message_idx !== ` +
|
|
@@ -237,11 +245,6 @@ class ChatContextManager {
|
|
|
237
245
|
this.startingLLMContextLength = undefined;
|
|
238
246
|
this.pendingMessages = undefined;
|
|
239
247
|
this.curAgentMsgIdx = undefined;
|
|
240
|
-
return {
|
|
241
|
-
type: "session_error",
|
|
242
|
-
session_id: this.sessionUUID,
|
|
243
|
-
message: errMsg,
|
|
244
|
-
};
|
|
245
248
|
}
|
|
246
249
|
processAgentMessage(msg, end) {
|
|
247
250
|
(0, assert_1.strict)(typeof this.startingLLMContextLength !== "undefined");
|
|
@@ -7,10 +7,9 @@ const errors_1 = require("../protocol/errors");
|
|
|
7
7
|
const logger = (0, sdk_1.getLogger)();
|
|
8
8
|
const HEARTBEAT_INTERVAL_MS = parseInt(process.env["HEARTBEAT_INTERVAL_MS"] || String(30 * 1000));
|
|
9
9
|
class Connection {
|
|
10
|
-
constructor(ws, userId,
|
|
10
|
+
constructor(ws, userId, heartbeat) {
|
|
11
11
|
this.ws = ws;
|
|
12
12
|
this.userId = userId;
|
|
13
|
-
this.apiKey = apiKey;
|
|
14
13
|
this.heartbeat = heartbeat;
|
|
15
14
|
this.lastMessageTime = Date.now();
|
|
16
15
|
}
|
|
@@ -25,16 +24,15 @@ class Connection {
|
|
|
25
24
|
*/
|
|
26
25
|
class ConnectionManager {
|
|
27
26
|
// this is a singleton object and can only be initialized by .init method
|
|
28
|
-
constructor(
|
|
29
|
-
this.apiKeyManager = apiKeyManager;
|
|
27
|
+
constructor(sessionMessageProcessor) {
|
|
30
28
|
this.sessionMessageProcessor = sessionMessageProcessor;
|
|
31
29
|
// Connections by id
|
|
32
30
|
this.connections = new Map();
|
|
33
31
|
// userId -> connectionIds
|
|
34
32
|
this.userToConnections = new Map();
|
|
35
33
|
}
|
|
36
|
-
static init(
|
|
37
|
-
const connManager = new ConnectionManager(
|
|
34
|
+
static init(createMsgProcessor) {
|
|
35
|
+
const connManager = new ConnectionManager({});
|
|
38
36
|
connManager.sessionMessageProcessor = createMsgProcessor(connManager);
|
|
39
37
|
return connManager;
|
|
40
38
|
}
|
|
@@ -42,25 +40,23 @@ class ConnectionManager {
|
|
|
42
40
|
* Handle new WebSocket connection with handshake protocol
|
|
43
41
|
* TODO: proper error code returning
|
|
44
42
|
*/
|
|
45
|
-
async handleConnection(webSocket,
|
|
43
|
+
async handleConnection(webSocket, token, req) {
|
|
46
44
|
const connectionId = (0, uuid_1.v4)();
|
|
47
45
|
logger.info(`[ConnectionManager] New connection: ${connectionId}`);
|
|
48
46
|
try {
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
if (!userData) {
|
|
52
|
-
logger.error(`[ConnectionManager] Invalid API key: ${apiKey}`);
|
|
47
|
+
const userUUID = await this.sessionMessageProcessor.authenticate(token);
|
|
48
|
+
if (!userUUID) {
|
|
53
49
|
throw new errors_1.ChatFatalError("invalid api key");
|
|
54
50
|
}
|
|
55
51
|
// Heartbeat interval
|
|
56
52
|
const heartbeat = setInterval(this.onHeartbeat.bind(this, connectionId), HEARTBEAT_INTERVAL_MS);
|
|
57
53
|
// Register connection
|
|
58
|
-
this.connections.set(connectionId, new Connection(webSocket,
|
|
54
|
+
this.connections.set(connectionId, new Connection(webSocket, userUUID, heartbeat));
|
|
59
55
|
// Add to user connections
|
|
60
|
-
if (!this.userToConnections.has(
|
|
61
|
-
this.userToConnections.set(
|
|
56
|
+
if (!this.userToConnections.has(userUUID)) {
|
|
57
|
+
this.userToConnections.set(userUUID, new Set());
|
|
62
58
|
}
|
|
63
|
-
const userConnections = this.userToConnections.get(
|
|
59
|
+
const userConnections = this.userToConnections.get(userUUID);
|
|
64
60
|
if (userConnections) {
|
|
65
61
|
userConnections.add(connectionId);
|
|
66
62
|
}
|
|
@@ -71,7 +67,7 @@ class ConnectionManager {
|
|
|
71
67
|
const response = {
|
|
72
68
|
t: "ready",
|
|
73
69
|
c_id: connectionId,
|
|
74
|
-
user_uuid:
|
|
70
|
+
user_uuid: userUUID,
|
|
75
71
|
};
|
|
76
72
|
// Check if connection still exists before sending
|
|
77
73
|
if (this.connections.has(connectionId)) {
|
|
@@ -82,7 +78,7 @@ class ConnectionManager {
|
|
|
82
78
|
` before ready message could be sent`);
|
|
83
79
|
}
|
|
84
80
|
logger.info(`[ConnectionManager] Connection ${connectionId} registered ` +
|
|
85
|
-
`for user ${
|
|
81
|
+
`for user ${userUUID}, version: ${clientVersion}`);
|
|
86
82
|
}
|
|
87
83
|
catch (error) {
|
|
88
84
|
const errorMessage = error instanceof Error ? error.message : String(error);
|
|
@@ -226,7 +222,7 @@ class ConnectionManager {
|
|
|
226
222
|
userConnections.delete(connectionId);
|
|
227
223
|
if (userConnections.size === 0) {
|
|
228
224
|
this.userToConnections.delete(conn.userId);
|
|
229
|
-
this.sessionMessageProcessor.handleUserDisconnect(conn.userId);
|
|
225
|
+
const _ = this.sessionMessageProcessor.handleUserDisconnect(conn.userId);
|
|
230
226
|
}
|
|
231
227
|
}
|
|
232
228
|
this.connections.delete(connectionId);
|
|
@@ -330,23 +326,5 @@ class ConnectionManager {
|
|
|
330
326
|
const connections = this.userToConnections.get(userId);
|
|
331
327
|
return connections ? Array.from(connections) : [];
|
|
332
328
|
}
|
|
333
|
-
/**
|
|
334
|
-
* Implementation of IUserConnectionManager interface.
|
|
335
|
-
* Get cached api key for a user. This key will not be available
|
|
336
|
-
* after the user disconnects.
|
|
337
|
-
*/
|
|
338
|
-
getLiveUserApiKey(userId) {
|
|
339
|
-
const connectionIds = this.getUserConnections(userId);
|
|
340
|
-
// Return API key from the first active connection
|
|
341
|
-
// (all connections for a user should have the same API key)
|
|
342
|
-
for (const connectionId of connectionIds) {
|
|
343
|
-
const conn = this.connections.get(connectionId);
|
|
344
|
-
if (conn) {
|
|
345
|
-
return conn.apiKey;
|
|
346
|
-
}
|
|
347
|
-
}
|
|
348
|
-
// No active connections for this user
|
|
349
|
-
return undefined;
|
|
350
|
-
}
|
|
351
329
|
}
|
|
352
330
|
exports.ConnectionManager = ConnectionManager;
|
|
@@ -27,7 +27,7 @@ vitest_1.vi.mock("ws");
|
|
|
27
27
|
// to access private constructor. Since the constructor is private,
|
|
28
28
|
// we need to use type assertion
|
|
29
29
|
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
|
30
|
-
connectionManager = new connectionManager_1.ConnectionManager(
|
|
30
|
+
connectionManager = new connectionManager_1.ConnectionManager(mockSessionRegistry.mock);
|
|
31
31
|
});
|
|
32
32
|
(0, vitest_1.afterEach)(() => {
|
|
33
33
|
vitest_1.vi.restoreAllMocks();
|
|
@@ -37,7 +37,7 @@ vitest_1.vi.mock("ws");
|
|
|
37
37
|
await connectionManager.handleConnection(mockWebSocket.mock, "valid-api-key", { headers: {} });
|
|
38
38
|
// Wait for deferred send
|
|
39
39
|
await new Promise((resolve) => setImmediate(resolve));
|
|
40
|
-
(0, vitest_1.expect)(
|
|
40
|
+
(0, vitest_1.expect)(mockSessionRegistry.spies.authenticate).toHaveBeenCalledWith("valid-api-key");
|
|
41
41
|
(0, vitest_1.expect)(mockWebSocket.spies.send).toHaveBeenCalledWith(vitest_1.expect.stringMatching(/"t":"ready"/));
|
|
42
42
|
});
|
|
43
43
|
(0, vitest_1.it)("should reject connection with invalid API key", async () => {
|
|
@@ -61,6 +61,7 @@ vitest_1.vi.mock("ws");
|
|
|
61
61
|
d: {
|
|
62
62
|
type: "user_msg",
|
|
63
63
|
user_uuid: "user_1",
|
|
64
|
+
user_nickname: "User 1",
|
|
64
65
|
session_id: "session_1",
|
|
65
66
|
message: "message text",
|
|
66
67
|
message_idx: 12,
|
|
@@ -72,14 +73,6 @@ vitest_1.vi.mock("ws");
|
|
|
72
73
|
connectionManager.sendToUsers(new Set([mockFactories_1.MOCK_USERS.owner.uuid]), testMessage.d);
|
|
73
74
|
(0, vitest_1.expect)(mockWebSocket.spies.send).toHaveBeenCalledWith(JSON.stringify(testMessage));
|
|
74
75
|
});
|
|
75
|
-
(0, vitest_1.it)("should return API key for connected user", () => {
|
|
76
|
-
const apiKey = connectionManager.getLiveUserApiKey(mockFactories_1.MOCK_USERS.owner.uuid);
|
|
77
|
-
(0, vitest_1.expect)(apiKey).toBe("valid-api-key");
|
|
78
|
-
});
|
|
79
|
-
(0, vitest_1.it)("should return undefined for non-existent user", () => {
|
|
80
|
-
const apiKey = connectionManager.getLiveUserApiKey(mockFactories_1.MOCK_USERS.nonExistent.uuid);
|
|
81
|
-
(0, vitest_1.expect)(apiKey).toBeUndefined();
|
|
82
|
-
});
|
|
83
76
|
});
|
|
84
77
|
(0, vitest_1.describe)("connection cleanup", () => {
|
|
85
78
|
(0, vitest_1.beforeEach)(async () => {
|
|
@@ -93,9 +86,6 @@ vitest_1.vi.mock("ws");
|
|
|
93
86
|
const closeHandler = closeHandlerCall[1];
|
|
94
87
|
closeHandler();
|
|
95
88
|
}
|
|
96
|
-
// After cleanup, API key should no longer be available
|
|
97
|
-
const apiKey = connectionManager.getLiveUserApiKey(mockFactories_1.MOCK_USERS.owner.uuid);
|
|
98
|
-
(0, vitest_1.expect)(apiKey).toBeUndefined();
|
|
99
89
|
});
|
|
100
90
|
(0, vitest_1.it)("should clean up connection state on error", () => {
|
|
101
91
|
// Get the error handler from mock calls
|
|
@@ -105,9 +95,6 @@ vitest_1.vi.mock("ws");
|
|
|
105
95
|
const errorHandler = errorHandlerCall[1];
|
|
106
96
|
errorHandler(new Error("WebSocket error"));
|
|
107
97
|
}
|
|
108
|
-
// After cleanup, API key should no longer be available
|
|
109
|
-
const apiKey = connectionManager.getLiveUserApiKey(mockFactories_1.MOCK_USERS.owner.uuid);
|
|
110
|
-
(0, vitest_1.expect)(apiKey).toBeUndefined();
|
|
111
98
|
});
|
|
112
99
|
});
|
|
113
100
|
(0, vitest_1.describe)("message routing", () => {
|
|
@@ -5,9 +5,11 @@ Object.defineProperty(exports, "__esModule", { value: true });
|
|
|
5
5
|
exports.MESSAGE_INDEX_SUB_INCREMENT = exports.MESSAGE_INDEX_FULL_INCREMENT = exports.MESSAGE_INDEX_START_VALUE = void 0;
|
|
6
6
|
exports.sessionMessagesToNextIndex = sessionMessagesToNextIndex;
|
|
7
7
|
exports.sessionMessagesToLLMConversation = sessionMessagesToLLMConversation;
|
|
8
|
-
exports.
|
|
8
|
+
exports.userMessageToConversationMessage = userMessageToConversationMessage;
|
|
9
|
+
exports.assistantMessageToConversationMessage = assistantMessageToConversationMessage;
|
|
10
|
+
exports.toolResultToConversationMessage = toolResultToConversationMessage;
|
|
11
|
+
exports.sessionMessagesToConversationMessages = sessionMessagesToConversationMessages;
|
|
9
12
|
exports.llmUserMessageToUserMessageData = llmUserMessageToUserMessageData;
|
|
10
|
-
exports.userMessageToChatMessage = userMessageToChatMessage;
|
|
11
13
|
exports.chatToolResultMessageToSessionMessage = chatToolResultMessageToSessionMessage;
|
|
12
14
|
exports.chatUserMessageToSessionMessage = chatUserMessageToSessionMessage;
|
|
13
15
|
exports.chatAgentMessageToSessionMessage = chatAgentMessageToSessionMessage;
|
|
@@ -70,36 +72,82 @@ function sessionMessagesToLLMConversation(sessionMessages) {
|
|
|
70
72
|
}
|
|
71
73
|
return { firstIndex, conversation };
|
|
72
74
|
}
|
|
73
|
-
|
|
75
|
+
/**
|
|
76
|
+
* Convert the DB data (SessionMessage) for a ChatCompletionUserMessageParam
|
|
77
|
+
* into a ServerUserMessage message.
|
|
78
|
+
*/
|
|
79
|
+
function userMessageToConversationMessage(userMessage, user_uuid, message_idx, session_id) {
|
|
80
|
+
// The name on the message should be the uuid
|
|
81
|
+
const userMsgData = llmUserMessageToUserMessageData(userMessage);
|
|
82
|
+
const user_nickname = userMessage.name;
|
|
83
|
+
(0, assert_1.strict)(user_nickname, `ChatCompletionUserMessageParam without user ${JSON.stringify(userMessage)}`);
|
|
84
|
+
const msg = {
|
|
85
|
+
type: "user_msg",
|
|
86
|
+
message: userMsgData.message,
|
|
87
|
+
message_idx,
|
|
88
|
+
user_uuid,
|
|
89
|
+
user_nickname,
|
|
90
|
+
session_id,
|
|
91
|
+
};
|
|
92
|
+
if (userMsgData.imageB64) {
|
|
93
|
+
msg.imageB64 = userMsgData.imageB64;
|
|
94
|
+
}
|
|
95
|
+
return msg;
|
|
96
|
+
}
|
|
97
|
+
/**
|
|
98
|
+
* Convert the DB data (SessionMessage) for a
|
|
99
|
+
* ChatCompletionAssistantMessageParam into a ServerAgentMessage message.
|
|
100
|
+
*/
|
|
101
|
+
function assistantMessageToConversationMessage(msg, message_idx, session_id) {
|
|
102
|
+
(0, assert_1.strict)(!msg.audio);
|
|
103
|
+
return {
|
|
104
|
+
type: "agent_msg",
|
|
105
|
+
message: msg,
|
|
106
|
+
message_idx,
|
|
107
|
+
session_id,
|
|
108
|
+
};
|
|
109
|
+
}
|
|
110
|
+
/**
|
|
111
|
+
* Convert the DB data (SessionMessage) for a ChatCompletionToolMessageParam
|
|
112
|
+
* into a ServerToolCallResult message.
|
|
113
|
+
*/
|
|
114
|
+
function toolResultToConversationMessage(toolCall, message_idx, session_id) {
|
|
115
|
+
return {
|
|
116
|
+
type: "tool_call_result",
|
|
117
|
+
result: toolCall,
|
|
118
|
+
message_idx,
|
|
119
|
+
session_id,
|
|
120
|
+
};
|
|
121
|
+
}
|
|
122
|
+
/**
|
|
123
|
+
* Convert a set of SessionMessages (stored in the DB) into Conversation
|
|
124
|
+
* messages (protocol which can be sent over the wire).
|
|
125
|
+
*/
|
|
126
|
+
function sessionMessagesToConversationMessages(sessionMessages, defaultUserUuid, session_id) {
|
|
74
127
|
const msgs = [];
|
|
75
128
|
for (const sm of sessionMessages) {
|
|
76
129
|
const ccmp = sm.content;
|
|
77
130
|
const message_idx = sm.message_idx;
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
131
|
+
const msgRole = ccmp.role;
|
|
132
|
+
switch (msgRole) {
|
|
133
|
+
case "system":
|
|
134
|
+
throw new Error("cannot convert system to ConversationMessage");
|
|
81
135
|
case "assistant":
|
|
82
|
-
(
|
|
83
|
-
if (ccmp.content) {
|
|
84
|
-
msgs.push({
|
|
85
|
-
type: "agent_msg",
|
|
86
|
-
message: ccmp,
|
|
87
|
-
message_idx,
|
|
88
|
-
session_id,
|
|
89
|
-
});
|
|
90
|
-
}
|
|
91
|
-
// TODO: do we want to convert tool calls etc?
|
|
136
|
+
msgs.push(assistantMessageToConversationMessage(ccmp, message_idx, session_id));
|
|
92
137
|
break;
|
|
93
138
|
case "user":
|
|
94
139
|
{
|
|
95
|
-
const
|
|
96
|
-
|
|
97
|
-
msgs.push(msg);
|
|
98
|
-
}
|
|
140
|
+
const user_uuid = sm.sender_uuid || defaultUserUuid;
|
|
141
|
+
msgs.push(userMessageToConversationMessage(ccmp, user_uuid, message_idx, session_id));
|
|
99
142
|
}
|
|
100
143
|
break;
|
|
101
|
-
|
|
144
|
+
case "tool":
|
|
145
|
+
msgs.push(toolResultToConversationMessage(ccmp, message_idx, session_id));
|
|
102
146
|
break;
|
|
147
|
+
default: {
|
|
148
|
+
const _ = ccmp;
|
|
149
|
+
throw new Error(`unexpected message role: ${msgRole}`);
|
|
150
|
+
}
|
|
103
151
|
}
|
|
104
152
|
}
|
|
105
153
|
return msgs;
|
|
@@ -122,11 +170,11 @@ function llmUserMessageToUserMessageData(userMessage) {
|
|
|
122
170
|
image = content.image_url.url;
|
|
123
171
|
break;
|
|
124
172
|
case "input_audio":
|
|
125
|
-
throw new errors_1.ChatErrorMessage("
|
|
173
|
+
throw new errors_1.ChatErrorMessage("llmUserMessageToUserMessageData: audio content not supported");
|
|
126
174
|
case "file":
|
|
127
|
-
throw new errors_1.ChatErrorMessage("
|
|
175
|
+
throw new errors_1.ChatErrorMessage("llmUserMessageToUserMessageData: file content not supported");
|
|
128
176
|
default:
|
|
129
|
-
throw new errors_1.ChatErrorMessage("
|
|
177
|
+
throw new errors_1.ChatErrorMessage("llmUserMessageToUserMessageData: unexpected content.type");
|
|
130
178
|
}
|
|
131
179
|
}
|
|
132
180
|
const finalMsg = {
|
|
@@ -137,25 +185,6 @@ function llmUserMessageToUserMessageData(userMessage) {
|
|
|
137
185
|
}
|
|
138
186
|
return finalMsg;
|
|
139
187
|
}
|
|
140
|
-
function userMessageToChatMessage(userMessage, message_idx, defaultUserUuid, session_id) {
|
|
141
|
-
// The name on the message should be the uuid
|
|
142
|
-
const userMsgData = llmUserMessageToUserMessageData(userMessage);
|
|
143
|
-
if (!userMsgData) {
|
|
144
|
-
return undefined;
|
|
145
|
-
}
|
|
146
|
-
const user_uuid = userMessage.name || defaultUserUuid;
|
|
147
|
-
const msg = {
|
|
148
|
-
type: "user_msg",
|
|
149
|
-
message: userMsgData.message,
|
|
150
|
-
message_idx,
|
|
151
|
-
user_uuid,
|
|
152
|
-
session_id,
|
|
153
|
-
};
|
|
154
|
-
if (userMsgData.imageB64) {
|
|
155
|
-
msg.imageB64 = userMsgData.imageB64;
|
|
156
|
-
}
|
|
157
|
-
return msg;
|
|
158
|
-
}
|
|
159
188
|
function chatToolResultMessageToSessionMessage(chatMessage) {
|
|
160
189
|
return {
|
|
161
190
|
message_idx: chatMessage.message_idx,
|
|
@@ -164,7 +193,7 @@ function chatToolResultMessageToSessionMessage(chatMessage) {
|
|
|
164
193
|
};
|
|
165
194
|
}
|
|
166
195
|
function chatUserMessageToSessionMessage(chatMessage) {
|
|
167
|
-
const userMsg = (0, agent_1.createUserMessage)(chatMessage.message, chatMessage.imageB64, chatMessage.
|
|
196
|
+
const userMsg = (0, agent_1.createUserMessage)(chatMessage.message, chatMessage.imageB64, chatMessage.user_nickname);
|
|
168
197
|
(0, assert_1.strict)(userMsg);
|
|
169
198
|
return {
|
|
170
199
|
message_idx: chatMessage.message_idx,
|