From 5c9787c14540190a4960edd483cf328da2097012 Mon Sep 17 00:00:00 2001
From: Mary Hipp <maryhipp@Marys-MacBook-Air.local>
Date: Wed, 2 Aug 2023 09:46:29 -0400
Subject: [PATCH] add project-id header to requests

---
 .../frontend/web/src/app/components/InvokeAIUI.tsx   | 11 +++++++++--
 invokeai/frontend/web/src/services/api/client.ts     | 12 ++++++++++--
 invokeai/frontend/web/src/services/api/index.ts      |  6 +++++-
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/invokeai/frontend/web/src/app/components/InvokeAIUI.tsx b/invokeai/frontend/web/src/app/components/InvokeAIUI.tsx
index 7df390bce6..d6e6c42728 100644
--- a/invokeai/frontend/web/src/app/components/InvokeAIUI.tsx
+++ b/invokeai/frontend/web/src/app/components/InvokeAIUI.tsx
@@ -13,7 +13,7 @@ import { addMiddleware, resetMiddlewares } from 'redux-dynamic-middlewares';
 import Loading from '../../common/components/Loading/Loading';
 
 import { Middleware } from '@reduxjs/toolkit';
-import { $authToken, $baseUrl } from 'services/api/client';
+import { $authToken, $baseUrl, $projectId } from 'services/api/client';
 import { socketMiddleware } from 'services/events/middleware';
 import '../../i18n';
 import { AddImageToBoardContextProvider } from '../contexts/AddImageToBoardContext';
@@ -37,6 +37,7 @@ const InvokeAIUI = ({
   config,
   headerComponent,
   middleware,
+  projectId,
 }: Props) => {
   useEffect(() => {
     // configure API client token
@@ -49,6 +50,11 @@ const InvokeAIUI = ({
       $baseUrl.set(apiUrl);
     }
 
+    // configure API client project header
+    if (apiUrl) {
+      $projectId.set(projectId);
+    }
+
     // reset dynamically added middlewares
     resetMiddlewares();
 
@@ -68,8 +74,9 @@ const InvokeAIUI = ({
       // Reset the API client token and base url on unmount
       $baseUrl.set(undefined);
       $authToken.set(undefined);
+      $projectId.set(undefined);
     };
-  }, [apiUrl, token, middleware]);
+  }, [apiUrl, token, middleware, projectId]);
 
   return (
     <React.StrictMode>
diff --git a/invokeai/frontend/web/src/services/api/client.ts b/invokeai/frontend/web/src/services/api/client.ts
index dd4caa460e..87deda7d36 100644
--- a/invokeai/frontend/web/src/services/api/client.ts
+++ b/invokeai/frontend/web/src/services/api/client.ts
@@ -16,6 +16,11 @@ export const $authToken = atom<string | undefined>();
  */
 export const $baseUrl = atom<string | undefined>();
 
+/**
+ * The optional project-id header.
+ */
+export const $projectId = atom<string | undefined>();
+
 /**
  * Autogenerated, type-safe fetch client for the API. Used when RTK Query is not an option.
  * Dynamically updates when the token or base url changes.
@@ -24,9 +29,12 @@ export const $baseUrl = atom<string | undefined>();
  * @example
  * const { get, post, del } = $client.get();
  */
-export const $client = computed([$authToken, $baseUrl], (authToken, baseUrl) =>
+export const $client = computed([$authToken, $baseUrl, $projectId], (authToken, baseUrl, projectId) =>
   createClient<paths>({
-    headers: authToken ? { Authorization: `Bearer ${authToken}` } : {},
+    headers: {
+      ...(authToken ? { Authorization: `Bearer ${authToken}` } : {}),
+      ...(projectId ? { "project-id": projectId } : {})
+    },
     // do not include `api/v1` in the base url for this client
     baseUrl: `${baseUrl ?? ''}`,
   })
diff --git a/invokeai/frontend/web/src/services/api/index.ts b/invokeai/frontend/web/src/services/api/index.ts
index 0a0391898c..a9de7130c9 100644
--- a/invokeai/frontend/web/src/services/api/index.ts
+++ b/invokeai/frontend/web/src/services/api/index.ts
@@ -6,7 +6,7 @@ import {
   createApi,
   fetchBaseQuery,
 } from '@reduxjs/toolkit/query/react';
-import { $authToken, $baseUrl } from 'services/api/client';
+import { $authToken, $baseUrl, $projectId } from 'services/api/client';
 
 export const tagTypes = [
   'Board',
@@ -30,6 +30,7 @@ const dynamicBaseQuery: BaseQueryFn<
 > = async (args, api, extraOptions) => {
   const baseUrl = $baseUrl.get();
   const authToken = $authToken.get();
+  const projectId = $projectId.get();
 
   const rawBaseQuery = fetchBaseQuery({
     baseUrl: `${baseUrl ?? ''}/api/v1`,
@@ -37,6 +38,9 @@ const dynamicBaseQuery: BaseQueryFn<
       if (authToken) {
         headers.set('Authorization', `Bearer ${authToken}`);
       }
+      if (projectId) {
+        headers.set("project-id", projectId)
+      }
 
       return headers;
     },