Skip to content

Commit

Permalink
[Obs AI Assistant] Use cookie auth for internal APIs in serverless te…
Browse files Browse the repository at this point in the history
…sts (elastic#203275)

## Summary

### Problem
Cookie authentication was introduced in Kibana for serverless internal
API tests via elastic#192727.
The serverless tests for Obs AI Assistant still uses API key based auth.

### Solution
Change authentication to cookie based auth for internal APIs in
serverless tests.

### Checklist

Check the PR satisfies following conditions. 

Reviewers should verify this PR satisfies this list as well.

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios
- [x] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)
  • Loading branch information
viduni94 authored Dec 9, 2024
1 parent fdedae0 commit 7eb0052
Show file tree
Hide file tree
Showing 14 changed files with 185 additions and 283 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -15,37 +15,47 @@ import supertest from 'supertest';
import { Subtract } from 'utility-types';
import { format } from 'url';
import { Config } from '@kbn/test';
import { SupertestWithRoleScope } from '@kbn/test-suites-xpack/api_integration/deployment_agnostic/services/role_scoped_supertest';
import { InheritedFtrProviderContext } from '../../../../services';
import type { InternalRequestHeader, RoleCredentials } from '../../../../../shared/services';

export function getObservabilityAIAssistantApiClient({
svlSharedConfig,
supertestUserWithCookieCredentials,
}: {
svlSharedConfig: Config;
supertestUserWithCookieCredentials?: SupertestWithRoleScope;
}) {
const kibanaServer = svlSharedConfig.get('servers.kibana');
const cAuthorities = svlSharedConfig.get('servers.kibana.certificateAuthorities');

const url = format({
...kibanaServer,
auth: false, // don't use auth in serverless
});

return createObservabilityAIAssistantApiClient(supertest.agent(url, { ca: cAuthorities }));
if (supertestUserWithCookieCredentials) {
return createObservabilityAIAssistantApiClient(supertestUserWithCookieCredentials);
} else {
const kibanaServer = svlSharedConfig.get('servers.kibana');
const cAuthorities = svlSharedConfig.get('servers.kibana.certificateAuthorities');

const url = format({
...kibanaServer,
auth: false, // don't use auth in serverless
});
return createObservabilityAIAssistantApiClient(supertest.agent(url, { ca: cAuthorities }));
}
}

type ObservabilityAIAssistantApiClientKey = 'slsUser';
type ObservabilityAIAssistantApiClientKey = 'slsAdmin' | 'slsEditor' | 'slsUser';

export type ObservabilityAIAssistantApiClient = Record<
ObservabilityAIAssistantApiClientKey,
Awaited<ReturnType<typeof getObservabilityAIAssistantApiClient>>
>;
export function createObservabilityAIAssistantApiClient(st: supertest.Agent) {

export function createObservabilityAIAssistantApiClient(
st: SupertestWithRoleScope | supertest.Agent
) {
return <TEndpoint extends ObservabilityAIAssistantAPIEndpoint>(
options: {
type?: 'form-data';
endpoint: TEndpoint;
roleAuthc: RoleCredentials;
internalReqHeader: InternalRequestHeader;
roleAuthc?: RoleCredentials;
internalReqHeader?: InternalRequestHeader;
} & ObservabilityAIAssistantAPIClientRequestParamsOf<TEndpoint> & {
params?: { query?: { _inspect?: boolean } };
}
Expand All @@ -57,7 +67,8 @@ export function createObservabilityAIAssistantApiClient(st: supertest.Agent) {
const { method, pathname, version } = formatRequest(endpoint, params.path);
const url = format({ pathname, query: params?.query });

const headers: Record<string, string> = { ...internalReqHeader, ...roleAuthc.apiKeyHeader };
const headers: Record<string, string> =
roleAuthc && internalReqHeader ? { ...internalReqHeader, ...roleAuthc.apiKeyHeader } : {};

if (version) {
headers['Elastic-Api-Version'] = version;
Expand Down Expand Up @@ -182,10 +193,34 @@ export async function getObservabilityAIAssistantApiClientService({
getService,
}: InheritedFtrProviderContext): Promise<ObservabilityAIAssistantApiClient> {
const svlSharedConfig = getService('config');
// defaults to elastic_admin user when used without auth
const roleScopedSupertest = getService('roleScopedSupertest');

const supertestAdminWithCookieCredentials: SupertestWithRoleScope =
await roleScopedSupertest.getSupertestWithRoleScope('admin', {
useCookieHeader: true,
withInternalHeaders: true,
});

const supertestEditorWithCookieCredentials: SupertestWithRoleScope =
await roleScopedSupertest.getSupertestWithRoleScope('editor', {
useCookieHeader: true,
withInternalHeaders: true,
});

return {
// defaults to elastic_admin user when used without auth
slsUser: await getObservabilityAIAssistantApiClient({
svlSharedConfig,
}),
// cookie auth for internal apis
slsAdmin: await getObservabilityAIAssistantApiClient({
svlSharedConfig,
supertestUserWithCookieCredentials: supertestAdminWithCookieCredentials,
}),
// cookie auth for internal apis
slsEditor: await getObservabilityAIAssistantApiClient({
svlSharedConfig,
supertestUserWithCookieCredentials: supertestEditorWithCookieCredentials,
}),
};
}
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import {
LlmProxy,
createLlmProxy,
} from '@kbn/test-suites-xpack/observability_ai_assistant_api_integration/common/create_llm_proxy';
import { SupertestWithRoleScope } from '@kbn/test-suites-xpack/api_integration/deployment_agnostic/services/role_scoped_supertest';
import { FtrProviderContext } from '../../common/ftr_provider_context';
import { createProxyActionConnector, deleteActionConnector } from '../../common/action_connectors';
import type { InternalRequestHeader, RoleCredentials } from '../../../../../../shared/services';
Expand All @@ -21,6 +22,9 @@ export default function ApiTest({ getService }: FtrProviderContext) {
const svlUserManager = getService('svlUserManager');
const svlCommonApi = getService('svlCommonApi');
const log = getService('log');
const roleScopedSupertest = getService('roleScopedSupertest');

let supertestEditorWithCookieCredentials: SupertestWithRoleScope;

const CHAT_API_URL = `/internal/observability_ai_assistant/chat`;

Expand Down Expand Up @@ -52,6 +56,15 @@ export default function ApiTest({ getService }: FtrProviderContext) {
before(async () => {
roleAuthc = await svlUserManager.createM2mApiKeyWithRoleScope('editor');
internalReqHeader = svlCommonApi.getInternalRequestHeader();

supertestEditorWithCookieCredentials = await roleScopedSupertest.getSupertestWithRoleScope(
'editor',
{
useCookieHeader: true,
withInternalHeaders: true,
}
);

proxy = await createLlmProxy(log);
connectorId = await createProxyActionConnector({
supertest: supertestWithoutAuth,
Expand All @@ -75,10 +88,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
});

it("returns a 4xx if the connector doesn't exist", async () => {
await supertestWithoutAuth
await supertestEditorWithCookieCredentials
.post(CHAT_API_URL)
.set(roleAuthc.apiKeyHeader)
.set(internalReqHeader)
.send({
name: 'my_api_call',
messages,
Expand All @@ -104,10 +115,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
const receivedChunks: Array<Record<string, any>> = [];

const passThrough = new PassThrough();
supertestWithoutAuth
supertestEditorWithCookieCredentials
.post(CHAT_API_URL)
.set(roleAuthc.apiKeyHeader)
.set(internalReqHeader)
.on('error', reject)
.send({
name: 'my_api_call',
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ import {
LlmResponseSimulator,
} from '@kbn/test-suites-xpack/observability_ai_assistant_api_integration/common/create_llm_proxy';
import { createOpenAiChunk } from '@kbn/test-suites-xpack/observability_ai_assistant_api_integration/common/create_openai_chunk';
import { SupertestWithRoleScope } from '@kbn/test-suites-xpack/api_integration/deployment_agnostic/services/role_scoped_supertest';
import { FtrProviderContext } from '../../common/ftr_provider_context';
import {
decodeEvents,
Expand All @@ -39,6 +40,9 @@ export default function ApiTest({ getService }: FtrProviderContext) {
const log = getService('log');
const svlUserManager = getService('svlUserManager');
const svlCommonApi = getService('svlCommonApi');
const roleScopedSupertest = getService('roleScopedSupertest');

let supertestEditorWithCookieCredentials: SupertestWithRoleScope;

const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantAPIClient');

Expand Down Expand Up @@ -82,10 +86,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
(body) => !isFunctionTitleRequest(body)
);
const responsePromise = new Promise<Response>((resolve, reject) => {
supertestWithoutAuth
supertestEditorWithCookieCredentials
.post(COMPLETE_API_URL)
.set(roleAuthc.apiKeyHeader)
.set(internalReqHeader)
.send({
messages,
connectorId,
Expand Down Expand Up @@ -134,6 +136,14 @@ export default function ApiTest({ getService }: FtrProviderContext) {
roleAuthc,
internalReqHeader,
});

supertestEditorWithCookieCredentials = await roleScopedSupertest.getSupertestWithRoleScope(
'editor',
{
useCookieHeader: true,
withInternalHeaders: true,
}
);
});

after(async () => {
Expand All @@ -155,10 +165,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {

const passThrough = new PassThrough();

supertestWithoutAuth
supertestEditorWithCookieCredentials
.post(COMPLETE_API_URL)
.set(roleAuthc.apiKeyHeader)
.set(internalReqHeader)
.send({
messages,
connectorId,
Expand Down Expand Up @@ -254,6 +262,7 @@ export default function ApiTest({ getService }: FtrProviderContext) {
},
});
});

describe('when creating a new conversation', () => {
let events: StreamingChatResponseEvent[];

Expand All @@ -273,12 +282,14 @@ export default function ApiTest({ getService }: FtrProviderContext) {
content: 'Hello',
},
});

expect(omit(events[1], 'id')).to.eql({
type: StreamingChatResponseEventType.ChatCompletionChunk,
message: {
content: ' again',
},
});

expect(omit(events[2], 'id', 'message.@timestamp')).to.eql({
type: StreamingChatResponseEventType.ChatCompletionMessage,
message: {
Expand Down Expand Up @@ -329,10 +340,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
)[0]?.conversation.id;

await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'DELETE /internal/observability_ai_assistant/conversation/{conversationId}',
roleAuthc,
internalReqHeader,
params: {
path: {
conversationId: createdConversationId,
Expand Down Expand Up @@ -417,10 +426,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
).to.eql(0);

const conversations = await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'POST /internal/observability_ai_assistant/conversations',
roleAuthc,
internalReqHeader,
})
.expect(200);

Expand Down Expand Up @@ -449,10 +456,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
.completeAfterIntercept();

const createResponse = await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'POST /internal/observability_ai_assistant/chat/complete',
roleAuthc,
internalReqHeader,
params: {
body: {
messages,
Expand All @@ -470,10 +475,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
conversationCreatedEvent = getConversationCreatedEvent(createResponse.body);

const conversationId = conversationCreatedEvent.conversation.id;
const fullConversation = await observabilityAIAssistantAPIClient.slsUser({
const fullConversation = await observabilityAIAssistantAPIClient.slsEditor({
endpoint: 'GET /internal/observability_ai_assistant/conversation/{conversationId}',
internalReqHeader,
roleAuthc,
params: {
path: {
conversationId,
Expand All @@ -486,10 +489,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
.completeAfterIntercept();

const updatedResponse = await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'POST /internal/observability_ai_assistant/chat/complete',
internalReqHeader,
roleAuthc,
params: {
body: {
messages: [
Expand Down Expand Up @@ -519,10 +520,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {

after(async () => {
await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'DELETE /internal/observability_ai_assistant/conversation/{conversationId}',
internalReqHeader,
roleAuthc,
params: {
path: {
conversationId: conversationCreatedEvent.conversation.id,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,6 @@ export default function ApiTest({ getService }: FtrProviderContext) {
const responseBody = await invokeChatCompleteWithFunctionRequest({
connectorId,
observabilityAIAssistantAPIClient,
internalReqHeader,
roleAuthc,
functionCall: {
name: ELASTICSEARCH_FUNCTION_NAME,
trigger: MessageRole.User,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ import {
} from '@kbn/observability-ai-assistant-plugin/common';
import type { AssistantScope } from '@kbn/ai-assistant-common';
import { Readable } from 'stream';
import type { InternalRequestHeader, RoleCredentials } from '../../../../../../../shared/services';
import { ObservabilityAIAssistantApiClient } from '../../../common/observability_ai_assistant_api_client';

function decodeEvents(body: Readable | string) {
Expand All @@ -34,22 +33,16 @@ export async function invokeChatCompleteWithFunctionRequest({
connectorId,
observabilityAIAssistantAPIClient,
functionCall,
roleAuthc,
internalReqHeader,
scopes,
}: {
connectorId: string;
observabilityAIAssistantAPIClient: ObservabilityAIAssistantApiClient;
functionCall: Message['message']['function_call'];
scopes?: AssistantScope[];
roleAuthc: RoleCredentials;
internalReqHeader: InternalRequestHeader;
}) {
const { body } = await observabilityAIAssistantAPIClient
.slsUser({
.slsEditor({
endpoint: 'POST /internal/observability_ai_assistant/chat/complete',
internalReqHeader,
roleAuthc,
params: {
body: {
messages: [
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,6 @@ export default function ApiTest({ getService }: FtrProviderContext) {
await invokeChatCompleteWithFunctionRequest({
connectorId,
observabilityAIAssistantAPIClient,
internalReqHeader,
roleAuthc,
functionCall: {
name: 'summarize',
trigger: MessageRole.User,
Expand All @@ -77,10 +75,8 @@ export default function ApiTest({ getService }: FtrProviderContext) {
});

it('persists entry in knowledge base', async () => {
const res = await observabilityAIAssistantAPIClient.slsUser({
const res = await observabilityAIAssistantAPIClient.slsEditor({
endpoint: 'GET /internal/observability_ai_assistant/kb/entries',
internalReqHeader,
roleAuthc,
params: {
query: {
query: '',
Expand Down
Loading

0 comments on commit 7eb0052

Please sign in to comment.