From f519520966d61204f6875d67af08cbd810e147bf Mon Sep 17 00:00:00 2001 From: Nawaz Dhandala Date: Wed, 24 Dec 2025 17:29:54 +0000 Subject: [PATCH] feat(ai-agent): Enhance AI Agent registration with name and description fields --- AIAgent/Config.ts | 6 +++ AIAgent/Services/Register.ts | 89 +++++++++++++++++++++++---------- Common/Server/API/AIAgentAPI.ts | 89 ++++++++++++++++++++++++++++++++- 3 files changed, 157 insertions(+), 27 deletions(-) diff --git a/AIAgent/Config.ts b/AIAgent/Config.ts index 3413143567..38b8719a09 100644 --- a/AIAgent/Config.ts +++ b/AIAgent/Config.ts @@ -23,6 +23,12 @@ if (!process.env["AI_AGENT_KEY"]) { export const AI_AGENT_KEY: string = process.env["AI_AGENT_KEY"]; +export const AI_AGENT_NAME: string | null = + process.env["AI_AGENT_NAME"] || null; + +export const AI_AGENT_DESCRIPTION: string | null = + process.env["AI_AGENT_DESCRIPTION"] || null; + export const HOSTNAME: string = process.env["HOSTNAME"] || "localhost"; export const PORT: Port = new Port( diff --git a/AIAgent/Services/Register.ts b/AIAgent/Services/Register.ts index 4c82c66b53..72bf9352cf 100644 --- a/AIAgent/Services/Register.ts +++ b/AIAgent/Services/Register.ts @@ -1,11 +1,19 @@ -import { ONEUPTIME_URL, AI_AGENT_ID, AI_AGENT_KEY } from "../Config"; +import { + ONEUPTIME_URL, + AI_AGENT_ID, + AI_AGENT_KEY, + AI_AGENT_NAME, + AI_AGENT_DESCRIPTION, +} from "../Config"; import HTTPResponse from "Common/Types/API/HTTPResponse"; import URL from "Common/Types/API/URL"; import { JSONObject } from "Common/Types/JSON"; import Sleep from "Common/Types/Sleep"; import API from "Common/Utils/API"; +import { HasClusterKey } from "Common/Server/EnvironmentConfig"; import LocalCache from "Common/Server/Infrastructure/LocalCache"; import logger from "Common/Server/Utils/Logger"; +import ClusterKeyAuthorization from "Common/Server/Middleware/ClusterKeyAuthorization"; export default class Register { public static async registerAIAgent(): Promise { @@ -35,36 +43,65 @@ export default class Register { } private static async _registerAIAgent(): Promise { - // Validate AI agent by sending alive request - if (!AI_AGENT_ID) { - logger.error("AI_AGENT_ID should be set"); - return process.exit(); - } + if (HasClusterKey) { + // Clustered mode: Auto-register and get ID from server + const aiAgentRegistrationUrl: URL = URL.fromString( + ONEUPTIME_URL.toString(), + ).addRoute("/api/ai-agent/register"); - const aliveUrl: URL = URL.fromString(ONEUPTIME_URL.toString()).addRoute( - "/api/ai-agent/alive", - ); + logger.debug("Registering AI Agent..."); + logger.debug("Sending request to: " + aiAgentRegistrationUrl.toString()); - logger.debug("Registering AI Agent..."); - logger.debug("Sending request to: " + aliveUrl.toString()); + const result: HTTPResponse = await API.post({ + url: aiAgentRegistrationUrl, + data: { + aiAgentKey: AI_AGENT_KEY, + aiAgentName: AI_AGENT_NAME, + aiAgentDescription: AI_AGENT_DESCRIPTION, + clusterKey: ClusterKeyAuthorization.getClusterKey(), + }, + }); - const result: HTTPResponse = await API.post({ - url: aliveUrl, - data: { - aiAgentKey: AI_AGENT_KEY.toString(), - aiAgentId: AI_AGENT_ID.toString(), - }, - }); + if (result.isSuccess()) { + logger.debug("AI Agent Registered"); + logger.debug(result.data); - if (result.isSuccess()) { - LocalCache.setString( - "AI_AGENT", - "AI_AGENT_ID", - AI_AGENT_ID.toString() as string, - ); - logger.debug("AI Agent registered successfully"); + const aiAgentId: string = result.data["_id"] as string; + + LocalCache.setString("AI_AGENT", "AI_AGENT_ID", aiAgentId as string); + } } else { - throw new Error("Failed to register AI Agent: " + result.statusCode); + // Non-clustered mode: Validate AI agent by sending alive request + if (!AI_AGENT_ID) { + logger.error("AI_AGENT_ID or ONEUPTIME_SECRET should be set"); + return process.exit(); + } + + const aliveUrl: URL = URL.fromString(ONEUPTIME_URL.toString()).addRoute( + "/api/ai-agent/alive", + ); + + logger.debug("Registering AI Agent..."); + logger.debug("Sending request to: " + aliveUrl.toString()); + + const result: HTTPResponse = await API.post({ + url: aliveUrl, + data: { + aiAgentKey: AI_AGENT_KEY.toString(), + aiAgentId: AI_AGENT_ID.toString(), + }, + }); + + if (result.isSuccess()) { + LocalCache.setString( + "AI_AGENT", + "AI_AGENT_ID", + AI_AGENT_ID.toString() as string, + ); + logger.debug("AI Agent registered successfully"); + } else { + throw new Error("Failed to register AI Agent: " + result.statusCode); + } } logger.debug( diff --git a/Common/Server/API/AIAgentAPI.ts b/Common/Server/API/AIAgentAPI.ts index 7b72dee688..3ddd5e863e 100644 --- a/Common/Server/API/AIAgentAPI.ts +++ b/Common/Server/API/AIAgentAPI.ts @@ -1,4 +1,5 @@ import UserMiddleware from "../Middleware/UserAuthorization"; +import ClusterKeyAuthorization from "../Middleware/ClusterKeyAuthorization"; import AIAgentService, { Service as AIAgentServiceType, } from "../Services/AIAgentService"; @@ -11,15 +12,101 @@ import Response from "../Utils/Response"; import BaseAPI from "./BaseAPI"; import LIMIT_MAX from "../../Types/Database/LimitMax"; import PositiveNumber from "../../Types/PositiveNumber"; -import AIAgent from "../../Models/DatabaseModels/AIAgent"; +import AIAgent, { + AIAgentConnectionStatus, +} from "../../Models/DatabaseModels/AIAgent"; import BadDataException from "../../Types/Exception/BadDataException"; import { JSONObject } from "../../Types/JSON"; import ObjectID from "../../Types/ObjectID"; +import OneUptimeDate from "../../Types/Date"; +import Version from "../../Types/Version"; export default class AIAgentAPI extends BaseAPI { public constructor() { super(AIAgent, AIAgentService); + // Register Global AI Agent. Custom AI Agent can be registered via dashboard. + this.router.post( + `${new this.entityType().getCrudApiPath()?.toString()}/register`, + ClusterKeyAuthorization.isAuthorizedServiceMiddleware, + async ( + req: ExpressRequest, + res: ExpressResponse, + next: NextFunction, + ): Promise => { + try { + const data: JSONObject = req.body; + + if (!data["aiAgentKey"]) { + return Response.sendErrorResponse( + req, + res, + new BadDataException("aiAgentKey is missing"), + ); + } + + const aiAgentKey: string = data["aiAgentKey"] as string; + + const aiAgent: AIAgent | null = await AIAgentService.findOneBy({ + query: { + key: aiAgentKey, + isGlobalAIAgent: true, + }, + select: { + _id: true, + }, + props: { + isRoot: true, + }, + }); + + if (aiAgent) { + await AIAgentService.updateOneById({ + id: aiAgent.id!, + data: { + name: (data["aiAgentName"] as string) || "Global AI Agent", + description: data["aiAgentDescription"] as string, + lastAlive: OneUptimeDate.getCurrentDate(), + connectionStatus: AIAgentConnectionStatus.Connected, + }, + props: { + isRoot: true, + }, + }); + + return Response.sendJsonObjectResponse(req, res, { + _id: aiAgent._id?.toString(), + message: "AI Agent already registered", + }); + } + + let newAIAgent: AIAgent = new AIAgent(); + newAIAgent.isGlobalAIAgent = true; + newAIAgent.key = aiAgentKey; + newAIAgent.name = + (data["aiAgentName"] as string) || "Global AI Agent"; + newAIAgent.description = data["aiAgentDescription"] as string; + newAIAgent.lastAlive = OneUptimeDate.getCurrentDate(); + newAIAgent.connectionStatus = AIAgentConnectionStatus.Connected; + newAIAgent.aiAgentVersion = new Version("1.0.0"); + + newAIAgent = await AIAgentService.create({ + data: newAIAgent, + props: { + isRoot: true, + }, + }); + + return Response.sendJsonObjectResponse(req, res, { + _id: newAIAgent._id?.toString(), + message: "AI Agent registered successfully", + }); + } catch (err) { + return next(err); + } + }, + ); + // Alive endpoint for AI Agent heartbeat this.router.post( `${new this.entityType().getCrudApiPath()?.toString()}/alive`,