Skip to content

Commit

Permalink
Collect request body bytes of the outgoing requests from the connectors
Browse files Browse the repository at this point in the history
  • Loading branch information
@ersin-erdal
ersin-erdal committed Jul 9, 2024
1 parent da28c23 commit 4d180d6
Show file tree
Hide file tree
Showing 108 changed files with 1,548 additions and 771 deletions.
45 changes: 44 additions & 1 deletion x-pack/plugins/actions/server/lib/action_executor.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ import {
} from '@kbn/core/server/mocks';
import { eventLoggerMock } from '@kbn/event-log-plugin/server/mocks';
import { spacesServiceMock } from '@kbn/spaces-plugin/server/spaces_service/spaces_service.mock';
import { ActionType as ConnectorType } from '../types';
import { ActionType as ConnectorType, ConnectorMetricsCollector } from '../types';
import { actionsAuthorizationMock, actionsMock } from '../mocks';
import {
asBackgroundTaskExecutionSource,
Expand All @@ -31,6 +31,8 @@ import { SecurityConnectorFeatureId } from '../../common';
import { TaskErrorSource } from '@kbn/task-manager-plugin/common';
import { createTaskRunError, getErrorSource } from '@kbn/task-manager-plugin/server/task_running';
import { GEN_AI_TOKEN_COUNT_EVENT } from './event_based_telemetry';
import { Event } from './create_action_event_log_record_object';
import { set } from '@kbn/safer-lodash-set';

const actionExecutor = new ActionExecutor({ isESOCanEncrypt: true });
const services = actionsMock.createServices();
Expand Down Expand Up @@ -163,6 +165,7 @@ const getBaseExecuteStartEventLogDoc = (unsecured: boolean) => {
},
id: CONNECTOR_ID,
name: '1',
type_id: 'test',
},
...(unsecured
? {}
Expand Down Expand Up @@ -211,6 +214,10 @@ const getBaseExecuteEventLogDoc = (unsecured: boolean) => {
};
};

const addConnectorMetrics = (event: Event, value: number) => {
set(event, 'kibana.action.execution.metrics.request_body_bytes', value);
};

beforeEach(() => {
jest.resetAllMocks();
jest.clearAllMocks();
Expand Down Expand Up @@ -280,13 +287,15 @@ describe('Action Executor', () => {
},
params: { foo: true },
logger: loggerMock,
connectorMetricsCollector: expect.any(ConnectorMetricsCollector),
});

expect(loggerMock.debug).toBeCalledWith('executing action test:1: 1');
expect(eventLogger.logEvent).toHaveBeenCalledTimes(2);

const execStartDoc = getBaseExecuteStartEventLogDoc(executeUnsecure);
const execDoc = getBaseExecuteEventLogDoc(executeUnsecure);
addConnectorMetrics(execDoc, 0);
expect(eventLogger.logEvent).toHaveBeenNthCalledWith(1, execStartDoc);
expect(eventLogger.logEvent).toHaveBeenNthCalledWith(2, execDoc);
});
Expand Down Expand Up @@ -353,13 +362,15 @@ describe('Action Executor', () => {
params: { foo: true },
logger: loggerMock,
source: executionSource.source,
connectorMetricsCollector: expect.any(ConnectorMetricsCollector),
});

expect(loggerMock.debug).toBeCalledWith('executing action test:1: 1');
expect(eventLogger.logEvent).toHaveBeenCalledTimes(2);

const execStartDoc = getBaseExecuteStartEventLogDoc(executeUnsecure);
const execDoc = getBaseExecuteEventLogDoc(executeUnsecure);
addConnectorMetrics(execDoc, 0);
expect(eventLogger.logEvent).toHaveBeenNthCalledWith(1, {
...execStartDoc,
kibana: {
Expand Down Expand Up @@ -431,13 +442,15 @@ describe('Action Executor', () => {
},
params: { foo: true },
logger: loggerMock,
connectorMetricsCollector: expect.any(ConnectorMetricsCollector),
});

expect(loggerMock.debug).toBeCalledWith('executing action test:preconfigured: Preconfigured');
expect(eventLogger.logEvent).toHaveBeenCalledTimes(2);

const execStartDoc = getBaseExecuteStartEventLogDoc(executeUnsecure);
const execDoc = getBaseExecuteEventLogDoc(executeUnsecure);
addConnectorMetrics(execDoc, 0);
expect(eventLogger.logEvent).toHaveBeenNthCalledWith(1, {
...execStartDoc,
kibana: {
Expand Down Expand Up @@ -513,6 +526,7 @@ describe('Action Executor', () => {
params: { foo: true },
logger: loggerMock,
request: {},
connectorMetricsCollector: expect.any(ConnectorMetricsCollector),
});
}

Expand All @@ -532,6 +546,7 @@ describe('Action Executor', () => {

const execStartDoc = getBaseExecuteStartEventLogDoc(executeUnsecure);
const execDoc = getBaseExecuteEventLogDoc(executeUnsecure);
addConnectorMetrics(execDoc, 0);
expect(eventLogger.logEvent).toHaveBeenNthCalledWith(1, {
...execStartDoc,
kibana: {
Expand All @@ -540,6 +555,7 @@ describe('Action Executor', () => {
...execStartDoc.kibana.action,
id: 'system-connector-.cases',
name: 'System action: .cases',
type_id: '.cases',
},
saved_objects: [
{
Expand Down Expand Up @@ -569,6 +585,7 @@ describe('Action Executor', () => {
...execDoc.kibana.action,
id: 'system-connector-.cases',
name: 'System action: .cases',
type_id: '.cases',
},
saved_objects: [
{
Expand Down Expand Up @@ -890,6 +907,7 @@ describe('Action Executor', () => {
},
params: { foo: true },
logger: loggerMock,
connectorMetricsCollector: expect.any(ConnectorMetricsCollector),
});
});

Expand Down Expand Up @@ -921,6 +939,7 @@ describe('Action Executor', () => {
params: { foo: true },
logger: loggerMock,
request: {},
connectorMetricsCollector: expect.any(ConnectorMetricsCollector),
});
});

Expand Down Expand Up @@ -989,13 +1008,15 @@ describe('Action Executor', () => {
},
params: { foo: true },
logger: loggerMock,
connectorMetricsCollector: expect.any(ConnectorMetricsCollector),
});

expect(loggerMock.debug).toBeCalledWith('executing action test:preconfigured: Preconfigured');
expect(eventLogger.logEvent).toHaveBeenCalledTimes(2);

const execStartDoc = getBaseExecuteStartEventLogDoc(executeUnsecure);
const execDoc = getBaseExecuteEventLogDoc(executeUnsecure);
addConnectorMetrics(execDoc, 0);
expect(eventLogger.logEvent).toHaveBeenNthCalledWith(1, {
...execStartDoc,
kibana: {
Expand Down Expand Up @@ -1026,6 +1047,12 @@ describe('Action Executor', () => {
...execDoc.kibana.action,
id: 'preconfigured',
name: 'Preconfigured',
execution: {
...execStartDoc.kibana.action.execution,
metrics: {
request_body_bytes: 0,
},
},
},
saved_objects: [
{
Expand Down Expand Up @@ -1074,6 +1101,7 @@ describe('Action Executor', () => {
params: { foo: true },
logger: loggerMock,
request: {},
connectorMetricsCollector: expect.any(ConnectorMetricsCollector),
});

expect(loggerMock.debug).toBeCalledWith(
Expand All @@ -1083,6 +1111,7 @@ describe('Action Executor', () => {

const execStartDoc = getBaseExecuteStartEventLogDoc(executeUnsecure);
const execDoc = getBaseExecuteEventLogDoc(executeUnsecure);
addConnectorMetrics(execDoc, 0);
expect(eventLogger.logEvent).toHaveBeenNthCalledWith(1, {
...execStartDoc,
kibana: {
Expand All @@ -1091,6 +1120,7 @@ describe('Action Executor', () => {
...execStartDoc.kibana.action,
id: 'system-connector-.cases',
name: 'System action: .cases',
type_id: '.cases',
},
saved_objects: [
{
Expand Down Expand Up @@ -1120,6 +1150,7 @@ describe('Action Executor', () => {
...execDoc.kibana.action,
id: 'system-connector-.cases',
name: 'System action: .cases',
type_id: '.cases',
},
saved_objects: [
{
Expand Down Expand Up @@ -1290,6 +1321,7 @@ describe('Action Executor', () => {
},
params: { foo: true },
logger: loggerMock,
connectorMetricsCollector: expect.any(ConnectorMetricsCollector),
});
}
});
Expand Down Expand Up @@ -1385,6 +1417,7 @@ describe('Event log', () => {
},
name: undefined,
id: 'action1',
type_id: 'test',
},
alert: {
rule: {
Expand Down Expand Up @@ -1430,6 +1463,7 @@ describe('Event log', () => {
},
name: 'action-1',
id: '1',
type_id: 'test',
},
alert: {
rule: {
Expand Down Expand Up @@ -1483,6 +1517,7 @@ describe('Event log', () => {
},
name: 'action-1',
id: '1',
type_id: 'test',
},
alert: {
rule: {
Expand Down Expand Up @@ -1559,9 +1594,13 @@ describe('Event log', () => {
gen_ai: {
usage: mockGenAi.usage,
},
metrics: {
request_body_bytes: 0,
},
},
name: 'action-1',
id: '1',
type_id: '.gen-ai',
},
alert: {
rule: {
Expand Down Expand Up @@ -1655,9 +1694,13 @@ describe('Event log', () => {
total_tokens: 35,
},
},
metrics: {
request_body_bytes: 0,
},
},
name: 'action-1',
id: '1',
type_id: '.gen-ai',
},
alert: {
rule: {
Expand Down
15 changes: 10 additions & 5 deletions x-pack/plugins/actions/server/lib/action_executor.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ import { IEventLogger, SAVED_OBJECT_REL_PRIMARY } from '@kbn/event-log-plugin/se
import { createTaskRunError, TaskErrorSource } from '@kbn/task-manager-plugin/server';
import { getErrorSource } from '@kbn/task-manager-plugin/server/task_running';
import { GEN_AI_TOKEN_COUNT_EVENT } from './event_based_telemetry';
import { ConnectorMetricsService } from './connector_metrics_service';
import { ConnectorMetricsCollector } from './connector_metrics_collector';
import { getGenAiTokenTracking, shouldTrackGenAiToken } from './gen_ai_token_tracking';
import {
validateConfig,
Expand Down Expand Up @@ -294,6 +294,7 @@ export class ActionExecutor {
actionExecutionId,
isInMemory: this.actionInfo.isInMemory,
...(source ? { source } : {}),
actionTypeId: this.actionInfo.actionTypeId,
});

eventLogger.logEvent(event);
Expand Down Expand Up @@ -390,7 +391,7 @@ export class ActionExecutor {
},
async (span) => {
const { actionTypeRegistry, analyticsService, eventLogger } = this.actionExecutorContext!;
const connectorMetricsService = new ConnectorMetricsService();
const connectorMetricsCollector = new ConnectorMetricsCollector();

const actionInfo = await this.getActionInfoInternal(actionId, namespace.namespace);

Expand Down Expand Up @@ -479,6 +480,7 @@ export class ActionExecutor {
actionExecutionId,
isInMemory: this.actionInfo.isInMemory,
...(source ? { source } : {}),
actionTypeId,
});

eventLogger.startTiming(event);
Expand Down Expand Up @@ -512,7 +514,7 @@ export class ActionExecutor {
logger,
source,
...(actionType.isSystemActionType ? { request } : {}),
connectorMetricsService,
connectorMetricsCollector,
});

if (rawResult && rawResult.status === 'error') {
Expand Down Expand Up @@ -541,8 +543,6 @@ export class ActionExecutor {
status: 'ok',
};

// console.log('body bytes ==> ', connectorMetricsService.getRequestBodyByte());

event.event = event.event || {};

const { error, ...resultWithoutError } = result;
Expand All @@ -553,6 +553,11 @@ export class ActionExecutor {
event.user = event.user || {};
event.user.name = currentUser?.username;
event.user.id = currentUser?.profile_uid;
set(
event,
'kibana.action.execution.metrics.request_body_bytes',
connectorMetricsCollector.getRequestBodyByte()
);

if (result.status === 'ok') {
span?.setOutcome('success');
Expand Down
22 changes: 11 additions & 11 deletions x-pack/plugins/actions/server/lib/axios_utils.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ import {
import { loggingSystemMock } from '@kbn/core/server/mocks';
import { actionsConfigMock } from '../actions_config.mock';
import { getCustomAgents } from './get_custom_agents';
import { ConnectorMetricsService } from './connector_metrics_service';
import { ConnectorMetricsCollector } from './connector_metrics_collector';

const TestUrl = 'https://elastic.co/foo/bar/baz';

Expand Down Expand Up @@ -80,7 +80,7 @@ describe('request', () => {
});
});

test('adds request body bytes from request header on a successful request when connectorMetricsService is provided', async () => {
test('adds request body bytes from request header on a successful request when connectorMetricsCollector is provided', async () => {
const contentLength = 12;
axiosMock.mockImplementation(() => ({
status: 200,
Expand All @@ -90,17 +90,17 @@ describe('request', () => {
headers: { 'Content-Length': contentLength },
},
}));
const connectorMetricsService = new ConnectorMetricsService();
const connectorMetricsCollector = new ConnectorMetricsCollector();
await request({
axios,
url: '/test',
logger,
data: { test: 12345 },
configurationUtilities,
connectorMetricsService,
connectorMetricsCollector,
});

expect(connectorMetricsService.getRequestBodyByte()).toBe(contentLength);
expect(connectorMetricsCollector.getRequestBodyByte()).toBe(contentLength);
});

test('adds request body bytes from request header on a failed', async () => {
Expand All @@ -111,23 +111,23 @@ describe('request', () => {
headers: { 'Content-Length': contentLength },
})
);
const connectorMetricsService = new ConnectorMetricsService();
const connectorMetricsCollector = new ConnectorMetricsCollector();

try {
await request({
axios,
url: '/test',
logger,
configurationUtilities,
connectorMetricsService,
connectorMetricsCollector,
});
} catch (e) {
expect(connectorMetricsService.getRequestBodyByte()).toBe(contentLength);
expect(connectorMetricsCollector.getRequestBodyByte()).toBe(contentLength);
}
});

test('adds request body bytes from data when request header does not exist', async () => {
const connectorMetricsService = new ConnectorMetricsService();
const connectorMetricsCollector = new ConnectorMetricsCollector();
const data = { test: 12345 };

await request({
Expand All @@ -136,10 +136,10 @@ describe('request', () => {
logger,
data,
configurationUtilities,
connectorMetricsService,
connectorMetricsCollector,
});

expect(connectorMetricsService.getRequestBodyByte()).toBe(
expect(connectorMetricsCollector.getRequestBodyByte()).toBe(
Buffer.byteLength(JSON.stringify(data), 'utf8')
);
});
Expand Down
Loading

0 comments on commit 4d180d6

Please sign in to comment.