@agentscope/studio
Version:
AgentScope Studio is a powerful local monitoring and visualization tool designed to provide real-time insights into your system's performance and behavior.
301 lines (300 loc) • 13.7 kB
JavaScript
;
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.SocketManager = void 0;
const socket_io_1 = require("socket.io");
const Run_1 = require("../dao/Run");
const trpc_1 = require("../../../shared/src/types/trpc");
const messageForm_1 = require("../../../shared/src/types/messageForm");
const InputRequest_1 = require("../dao/InputRequest");
const ModelInvocation_1 = require("../dao/ModelInvocation");
class SocketManager {
static close() {
if (this.io) {
this.io.close();
}
}
static init(httpServer) {
this.io = new socket_io_1.Server(httpServer, {
cors: {
origin: '*',
},
});
// Python client connection
const pythonNamespace = this.io.of('/python');
pythonNamespace.on('connection', (socket) => {
const runId = socket.handshake.auth.run_id;
console.debug(`${socket.id}: Python client connected`);
socket.on('disconnect', () => __awaiter(this, void 0, void 0, function* () {
// Delete the input requests
yield InputRequest_1.InputRequestDao.deleteInputRequestsByRunId(runId);
this.changeRunStatusAndTriggerEvents(runId, messageForm_1.Status.DONE).catch((error) => {
console.error(error);
throw error;
});
}));
});
const clientNamespace = this.io.of('/client');
clientNamespace.on('connection', (socket) => {
console.debug('Client connected');
socket.on(trpc_1.SocketEvents.client.joinProjectListRoom, () => {
socket.join(trpc_1.SocketRoomName.ProjectListRoom);
console.debug(`${socket.id}: joined room: ${trpc_1.SocketRoomName.ProjectListRoom}`);
Run_1.RunDao.getAllProjects()
.then((projects) => {
// Push projects to the client
socket.emit(trpc_1.SocketEvents.server.pushProjects, projects);
})
.catch((error) => {
console.error(error);
throw error;
});
});
socket.on(trpc_1.SocketEvents.client.joinProjectRoom, (project, callback) => __awaiter(this, void 0, void 0, function* () {
const projectExist = yield Run_1.RunDao.doesProjectExist(project);
if (!projectExist) {
callback({
success: false,
message: `Project ${project} not found`,
});
}
else {
socket.join(`project-${project}`);
console.debug(`${socket.id}: joined room: project-${project}`);
// Return runs to this socket/client
Run_1.RunDao.getAllProjectRuns(project)
.then((runs) => {
// Push runs to the client
socket.emit(trpc_1.SocketEvents.server.pushRunsData, runs);
})
.catch((error) => {
console.error(error);
throw error;
});
}
}));
socket.on(trpc_1.SocketEvents.client.joinRunRoom, (runId, callback) => __awaiter(this, void 0, void 0, function* () {
const runExist = yield Run_1.RunDao.doesRunExist(runId);
if (!runExist) {
callback({
success: false,
message: `Run ${runId} not found`,
});
}
else {
socket.join(`run-${runId}`);
console.debug(`${socket.id}: joined room: run-${runId}`);
// Return run data, input requests and messages to this socket/client
Run_1.RunDao.getRunData(runId)
.then((data) => {
socket.emit(trpc_1.SocketEvents.server.pushRunData, data.runData);
socket.emit(trpc_1.SocketEvents.server.pushInputRequests, data.inputRequests);
socket.emit(trpc_1.SocketEvents.server.pushMessages, data.messages);
})
.catch((error) => {
console.error(error);
throw error;
});
// Return model invocation data
ModelInvocation_1.ModelInvocationDao.getModeInvocationData(runId).then((data) => {
socket.emit(trpc_1.SocketEvents.server.pushModelInvocationData, data);
});
}
}));
socket.on(trpc_1.SocketEvents.client.sendUserInputToServer, (requestId, blocksInput, structuredInput, callback) => __awaiter(this, void 0, void 0, function* () {
const inputRequest = yield InputRequest_1.InputRequestDao.getInputRequestByRequestId(requestId);
if (inputRequest === null) {
callback({
success: false,
message: `Input request ${requestId} not found`,
});
}
else {
const runId = inputRequest.runId;
yield InputRequest_1.InputRequestDao.deleteInputRequest(requestId);
// If input requests are empty, change the run status to finished
const res = yield Run_1.RunDao.getRunData(runId);
if (res.inputRequests.length === 0) {
this.changeRunStatusAndTriggerEvents(runId, messageForm_1.Status.RUNNING).catch((error) => {
console.error(error);
throw error;
});
}
// Emit the input to the python client
this.io
.of('/python')
.emit(trpc_1.SocketEvents.server.forwardUserInput, requestId, blocksInput, structuredInput);
}
}));
socket.on(trpc_1.SocketEvents.client.joinOverviewRoom, () => __awaiter(this, void 0, void 0, function* () {
socket.join(trpc_1.SocketRoomName.OverviewRoom);
console.debug(`${socket.id}: joined room: ${trpc_1.SocketRoomName.OverviewRoom}`);
// Return current overview data
const res = yield this._getOverViewData();
socket.emit(trpc_1.SocketEvents.server.pushOverviewData, res);
}));
socket.on(trpc_1.SocketEvents.client.leaveRoom, (room) => {
socket.leave(room);
console.debug(`${socket.id}: left room: ${room}`);
});
socket.on(trpc_1.SocketEvents.client.deleteProjects, (projects, callback) => __awaiter(this, void 0, void 0, function* () {
try {
yield Run_1.RunDao.deleteProjects(projects);
callback({
success: true,
message: `Success: ${projects.length} project deleted`,
});
// Update projectListRoom, overviewRoom,
this.broadcastOverviewDataToDashboardRoom();
this.broadcastRunToProjectListRoom();
}
catch (error) {
callback({
success: false,
message: `Error: ${error}`,
});
}
}));
socket.on(trpc_1.SocketEvents.client.deleteRuns, (runIds, callback) => __awaiter(this, void 0, void 0, function* () {
try {
const nDelete = yield Run_1.RunDao.deleteRuns(runIds);
callback({
success: nDelete === runIds.length,
message: `Deleted ${nDelete} runs`,
});
// Update data to overviewRoom, projectRoom
this.broadcastOverviewDataToDashboardRoom();
this.broadcastRunToProjectListRoom();
}
catch (error) {
callback({
success: false,
message: `Failed to delete runs: ${error}`,
});
}
}));
socket.on('disconnect', () => {
console.debug('Client disconnected');
});
});
}
/*
* Emit events to the project list room.
*/
static broadcastRunToProjectListRoom() {
Run_1.RunDao.getAllProjects()
.then((projects) => {
// Push projects to the client
this.io
.of('/client')
.to(trpc_1.SocketRoomName.ProjectListRoom)
.emit(trpc_1.SocketEvents.server.pushProjects, projects);
})
.catch((error) => {
console.error(error);
throw error;
});
}
static broadcastRunToProjectRoom(project) {
Run_1.RunDao.getAllProjectRuns(project)
.then((runs) => {
// Push runs to the client
this.io
.of('/client')
.to(`project-${project}`)
.emit(trpc_1.SocketEvents.server.pushRunsData, runs);
})
.catch((error) => {
console.error(error);
throw error;
});
}
/*
* Emit events to the run room.
*/
static broadcastMessageToRunRoom(runId, msgForm) {
this.io
.of('/client')
.to(`run-${runId}`)
.emit(trpc_1.SocketEvents.server.pushMessages, [
Object.assign({ id: msgForm.id, runId: msgForm.runId, replyId: msgForm.replyId }, msgForm.msg),
]);
}
static broadcastInputRequestToRunRoom(runId, inputRequest) {
this.io
.of('/client')
.to(`run-${runId}`)
.emit(trpc_1.SocketEvents.server.pushInputRequests, [inputRequest]);
}
static broadcastRunDataToRunRoom(runId, runData) {
this.io
.of('/client')
.to(`run-${runId}`)
.emit(trpc_1.SocketEvents.server.pushRunData, runData);
}
static clearInputRequestsToRunRoom(runId) {
this.io
.of('/client')
.to(`run-${runId}`)
.emit(trpc_1.SocketEvents.server.clearInputRequests);
}
static changeRunStatusAndTriggerEvents(runId, newStatus) {
return __awaiter(this, void 0, void 0, function* () {
const runExist = yield Run_1.RunDao.doesRunExist(runId);
if (runExist) {
// Update the run status to "finished"
yield Run_1.RunDao.changeRunStatus(runId, newStatus);
// Find the project by runId
const res = yield Run_1.RunDao.getRunData(runId);
const project = res.runData.project;
// Broadcast projects to all clients in the ProjectList room
this.broadcastRunToProjectListRoom();
// Broadcast runs to all clients in the project room
this.broadcastRunToProjectRoom(project);
// Broadcast run data to all clients in the run room
this.broadcastRunDataToRunRoom(runId, res.runData);
if (newStatus === messageForm_1.Status.DONE) {
// Clear the input requests for all clients in the run room
this.clearInputRequestsToRunRoom(runId);
}
}
});
}
static broadcastOverviewDataToDashboardRoom() {
this._getOverViewData()
.then((res) => {
this.io
.of('/client')
.to(trpc_1.SocketRoomName.OverviewRoom)
.emit(trpc_1.SocketEvents.server.pushOverviewData, res);
})
.catch((error) => {
console.error(error);
throw error;
});
}
static _getOverViewData() {
return __awaiter(this, void 0, void 0, function* () {
const res1 = yield Run_1.RunDao.getRunViewData();
const res2 = yield ModelInvocation_1.ModelInvocationDao.getModelInvocationViewData();
return Object.assign(Object.assign({}, res1), res2);
});
}
static broadcastModelInvocationDataToRunRoom(runId) {
ModelInvocation_1.ModelInvocationDao.getModeInvocationData(runId).then((data) => {
this.io
.of('/client')
.to(`run-${runId}`)
.emit(trpc_1.SocketEvents.server.pushModelInvocationData, data);
});
}
}
exports.SocketManager = SocketManager;