diff --git a/cli/commands/run/server/agent.controller.d.ts b/cli/commands/run/server/agent.controller.d.ts index 8e8197f..b75e90b 100644 --- a/cli/commands/run/server/agent.controller.d.ts +++ b/cli/commands/run/server/agent.controller.d.ts @@ -1,11 +1,12 @@ import { GetAgentHandlers } from "../../../utils/get.agent.handlers"; -export = AgentController +export = AgentController; declare class AgentController { - constructor(getAgentHandlers: GetAgentHandlers) - Initialize(call: any, callback: any): void - EvaluateBlock(call: any, callback: any): void - EvaluateTx(call: any, callback: any): void - EvaluateAlert(call: any, callback: any): void -} \ No newline at end of file + constructor(getAgentHandlers: GetAgentHandlers); + Initialize(call: any, callback: any): void; + EvaluateBlock(call: any, callback: any): void; + EvaluateTx(call: any, callback: any): void; + EvaluateAlert(call: any, callback: any): void; + HealthCheck(call: any, callback: any): void; +} diff --git a/cli/commands/run/server/agent.controller.js b/cli/commands/run/server/agent.controller.js index ea42967..a9f9a60 100644 --- a/cli/commands/run/server/agent.controller.js +++ b/cli/commands/run/server/agent.controller.js @@ -20,6 +20,31 @@ module.exports = class AgentController { this.initializeResponse = {}; } + async HealthCheck(call, callback) { + let errors = []; + let status = "SUCCESS"; + + if (this.healthCheck) { + try { + const response = await this.healthCheck(); + if (response && response.length > 0) { + status = "ERROR"; + errors = response.map((r) => ({ message: r })); + } + } catch (e) { + console.log(`${new Date().toISOString()} healthCheck`); + console.log(e); + status = "ERROR"; + errors = [{ message: e.message }]; + } + } + + callback(null, { + status, + errors, + }); + } + async Initialize(call, callback) { let status = "SUCCESS"; @@ -35,7 +60,7 @@ module.exports = class AgentController { } callback(null, { - status: status, + status, alertConfig: this.initializeResponse ? this.initializeResponse.alertConfig : undefined, @@ -152,6 +177,7 @@ module.exports = class AgentController { this.handleBlock = agentHandlers.handleBlock; this.handleTransaction = agentHandlers.handleTransaction; this.handleAlert = agentHandlers.handleAlert; + this.healthCheck = agentHandlers.healthCheck; } catch (e) { console.log(e); } diff --git a/cli/utils/get.agent.handlers.ts b/cli/utils/get.agent.handlers.ts index b3654ee..57b742b 100644 --- a/cli/utils/get.agent.handlers.ts +++ b/cli/utils/get.agent.handlers.ts @@ -1,10 +1,11 @@ -import { HandleBlock, HandleTransaction, HandleAlert, Initialize, InitializeResponse } from "../../sdk" +import { HandleBlock, HandleTransaction, HandleAlert, Initialize, InitializeResponse, HealthCheck } from "../../sdk" import { assertExists, assertIsNonEmptyString } from "." import { GetPythonAgentHandlers } from './get.python.agent.handlers' type AgentHandlers = { initialize?: Initialize, initializeResponse?: InitializeResponse | void, + healthCheck?: HealthCheck, handleTransaction?: HandleTransaction, handleBlock?: HandleBlock, handleAlert?: HandleAlert, diff --git a/sdk/handlers.ts b/sdk/handlers.ts index 2e2cc6f..f577d8b 100644 --- a/sdk/handlers.ts +++ b/sdk/handlers.ts @@ -5,6 +5,7 @@ import { InitializeResponse } from "./initialize.response" import { TransactionEvent } from "./transaction.event" export type Initialize = () => Promise +export type HealthCheck = () => Promise export type HandleTransaction = (txEvent: TransactionEvent) => Promise export type HandleBlock = (blockEvent: BlockEvent) => Promise export type HandleAlert = (alertEvent: AlertEvent) => Promise \ No newline at end of file diff --git a/sdk/index.ts b/sdk/index.ts index fa975f7..ffc6ad1 100644 --- a/sdk/index.ts +++ b/sdk/index.ts @@ -59,6 +59,7 @@ import { HandleTransaction, HandleBlock, HandleAlert, + HealthCheck } from "./handlers"; interface DiContainer { @@ -77,6 +78,7 @@ export { HandleTransaction, HandleBlock, HandleAlert, + HealthCheck, Finding, FindingSeverity, FindingType,