From 6c4ababc3fa147449bc31fe809ce31d889818026 Mon Sep 17 00:00:00 2001
From: flakey5 <73616808+flakey5@users.noreply.github.com>
Date: Tue, 14 May 2024 10:47:11 -0700
Subject: [PATCH] Add Authorization header to OpenAPI docs

Closes #43

Signed-off-by: flakey5 <73616808+flakey5@users.noreply.github.com>
---
 config.d.ts           | 10 +++-------
 index.ts              | 22 ++++++++++++++++++++++
 plugins/api.ts        |  7 +++++++
 schema.json           |  2 +-
 tests/e2e/api.test.ts | 41 +++++++++++++++++++++++++++++++++++++++++
 5 files changed, 74 insertions(+), 8 deletions(-)

diff --git a/config.d.ts b/config.d.ts
index e9f653d..69ab576 100644
--- a/config.d.ts
+++ b/config.d.ts
@@ -217,13 +217,9 @@ export interface AiWarpConfig {
         prefix?: string;
         path?: string;
       };
-      plugins?:
-        | {
-            [k: string]: unknown;
-          }
-        | {
-            [k: string]: unknown;
-          };
+      plugins?: {
+        [k: string]: unknown;
+      };
     }[];
   };
   module?: string;
diff --git a/index.ts b/index.ts
index ab02bd1..6ae4fcf 100644
--- a/index.ts
+++ b/index.ts
@@ -22,6 +22,28 @@ const stackable: Stackable<AiWarpConfig> = async function (fastify, opts) {
     })
   }
 
+  if (config.service === undefined) {
+    config.service = {}
+  }
+
+  const currentOpenApiConfig = typeof config.service.openapi === 'object' ? config.service.openapi : {}
+  if (config.auth?.jwt !== undefined) {
+    config.service.openapi = {
+      ...currentOpenApiConfig,
+      components: {
+        ...currentOpenApiConfig.components,
+        securitySchemes: {
+          ...currentOpenApiConfig.components?.securitySchemes,
+          aiWarpJwtToken: {
+            type: 'apiKey',
+            in: 'header',
+            name: 'Authorization'
+          }
+        }
+      }
+    }
+  }
+
   await fastify.register(platformaticService, opts)
 
   await fastify.register(warpPlugin, opts) // needs to be registered here for fastify.ai to be decorated
diff --git a/plugins/api.ts b/plugins/api.ts
index 4090ce1..0cb6c4e 100644
--- a/plugins/api.ts
+++ b/plugins/api.ts
@@ -11,10 +11,16 @@ function isAFastifyError (object: object): object is FastifyError {
 const InternalServerError = createError('INTERNAL_SERVER_ERROR', 'Internal Server Error', 500)
 
 const plugin: FastifyPluginAsyncTypebox = async (fastify) => {
+  let security: Array<Record<string, string[]>> | undefined
+  if (fastify.platformatic.config.auth?.jwt !== undefined) {
+    security = [{ aiWarpJwtToken: [] }]
+  }
+
   fastify.route({
     url: '/api/v1/prompt',
     method: 'POST',
     schema: {
+      security,
       body: Type.Object({
         prompt: Type.String(),
         chatHistory: Type.Optional(Type.Array(Type.Object({
@@ -55,6 +61,7 @@ const plugin: FastifyPluginAsyncTypebox = async (fastify) => {
     method: 'POST',
     schema: {
       produces: ['text/event-stream'],
+      security,
       body: Type.Object({
         prompt: Type.String(),
         chatHistory: Type.Optional(Type.Array(Type.Object({
diff --git a/schema.json b/schema.json
index e43822d..a34b6de 100644
--- a/schema.json
+++ b/schema.json
@@ -1,6 +1,6 @@
 {
   "$id": "ai-warp",
-  "version": "1.37.1",
+  "version": "1.39.0",
   "title": "Ai Warp Config",
   "type": "object",
   "properties": {
diff --git a/tests/e2e/api.test.ts b/tests/e2e/api.test.ts
index 99a6b90..c5b982f 100644
--- a/tests/e2e/api.test.ts
+++ b/tests/e2e/api.test.ts
@@ -197,6 +197,47 @@ it('provides all paths in OpenAPI', async () => {
   await app.close()
 })
 
+it('provides header for jwt token in OpenAPI', async () => {
+  const [app, port] = await buildAiWarpApp({
+    aiProvider: {
+      openai: {
+        model: 'gpt-3.5-turbo',
+        apiKey: ''
+      }
+    },
+    auth: {
+      jwt: {
+        secret: 'asd'
+      }
+    }
+  })
+
+  await app.start()
+
+  const res = await fetch(`http://localhost:${port}/documentation/json`)
+  const body = await res.json()
+
+  assert.deepStrictEqual(body.components, {
+    securitySchemes: {
+      aiWarpJwtToken: {
+        type: 'apiKey',
+        in: 'header',
+        name: 'Authorization'
+      }
+    },
+    schemas: {}
+  })
+
+  const endpoints = ['/api/v1/prompt', '/api/v1/stream']
+  for (const endpoint of endpoints) {
+    assert.deepStrictEqual(body.paths[endpoint].post.security, [
+      { aiWarpJwtToken: [] }
+    ])
+  }
+
+  await app.close()
+})
+
 it('prompt with wrong JSON', async () => {
   const [app, port] = await buildAiWarpApp({
     aiProvider: {