fix: missing messages, plan button, sources icons (#148)

Signed-off-by: Petr Kadlec <[email protected]>
i-am-bee · Dec 16, 2024 · a312c17 · a312c17
1 parent 5a1e34b
commit a312c17
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 11 deletions.
diff --git a/src/app/auth/rsc.tsx b/src/app/auth/rsc.tsx
@@ -21,6 +21,8 @@ import { cache } from 'react';
 import { SIGN_IN_PAGE, auth } from '.';
 import { addDaysToDate } from '@/utils/dates';
 import { readUser } from '../api/users';
+import { decodeMetadata } from '../api/utils';
+import { UserMetadata } from '@/store/user-profile/types';
 
 const DUMMY_JWT_TOKEN = process.env.DUMMY_JWT_TOKEN!;
 
@@ -31,15 +33,16 @@ export const getSession = cache(async () => {
 export const ensureSession = async () => {
   if (DUMMY_JWT_TOKEN) {
     const user = await readUser(DUMMY_JWT_TOKEN);
+
     if (user)
       return {
         user: {
           access_token: DUMMY_JWT_TOKEN,
         },
         expires: addDaysToDate(new Date(), SESSION_TEST_EXPIRY_DAYS),
         userProfile: {
-          metadata: {},
           ...user,
+          metadata: decodeMetadata<UserMetadata>(user.metadata),
           name: user.name ?? '',
           email: user.email ?? '',
           firstName: 'Test',

diff --git a/src/middleware.ts b/src/middleware.ts
@@ -90,7 +90,7 @@ const getCspHeader = (nonce: string) => {
       process.env.NODE_ENV === 'production' ? '' : `'unsafe-eval'`
     };
     style-src 'self' 'nonce-${nonce}';
-    img-src 'self' blob: data: www.ibm.com/;
+    img-src *;
     font-src 'self';
     object-src 'none';
     frame-src ${USERCONTENT_SITE_URL};

diff --git a/src/modules/chat/assistant-plan/PlanWithSources.tsx b/src/modules/chat/assistant-plan/PlanWithSources.tsx
@@ -42,6 +42,7 @@ import { TraceDataProvider } from '../trace/TraceDataProvider';
 import { FeatureName, isFeatureEnabled } from '@/utils/isFeatureEnabled';
 import { Spinner } from '@/components/Spinner/Spinner';
 import { useProjectContext } from '@/layout/providers/ProjectProvider';
+import { MAX_API_FETCH_LIMIT } from '@/app/api/utils';
 
 interface Props {
   message: BotChatMessage;
@@ -74,9 +75,7 @@ function PlanWithSourcesComponent({ message, inView }: Props) {
       project.id,
       thread?.id ?? '',
       message.run_id ?? '',
-      {
-        limit: 100,
-      },
+      PLAN_STEPS_QUERY_PARAMS,
     ),
     enabled: Boolean(!message.plan && thread && message.run_id && inView),
   });
@@ -234,3 +233,5 @@ const getSourcesWithSteps = (
 
   return uniqueSourcesWithSteps;
 };
+
+export const PLAN_STEPS_QUERY_PARAMS = { limit: MAX_API_FETCH_LIMIT };
diff --git a/src/modules/chat/providers/ChatProvider.tsx b/src/modules/chat/providers/ChatProvider.tsx
@@ -67,7 +67,12 @@ import { THREAD_TITLE_MAX_LENGTH } from '../history/ThreadItem';
 import { useGetThreadAssistant } from '../history/useGetThreadAssistant';
 import { useChatStream } from '../hooks/useChatStream';
 import { useThreadApi } from '../hooks/useThreadApi';
-import { messagesWithFilesQuery, readRunQuery, runsQuery } from '../queries';
+import {
+  messagesWithFilesQuery,
+  readRunQuery,
+  runsQuery,
+  runStepsQuery,
+} from '../queries';
 import {
   ChatMessage,
   MessageWithFiles,
@@ -87,6 +92,7 @@ import { AssistantBuilderState } from '@/modules/assistants/builder/Builder';
 import { useModal } from '@/layout/providers/ModalProvider';
 import { ApiError } from '@/app/api/errors';
 import { UsageLimitModal } from '@/components/UsageLimitModal/UsageLimitModal';
+import { PLAN_STEPS_QUERY_PARAMS } from '../assistant-plan/PlanWithSources';
 
 interface CancelRunParams {
   threadId: string;
@@ -195,7 +201,11 @@ export function ChatProvider({
   }, [setVectorStoreId, thread?.tool_resources?.file_search?.vector_store_ids]);
 
   const setMessagesWithFilesQueryData = useCallback(
-    (threadId?: string, newMessage?: MessageWithFiles | null) => {
+    (
+      threadId?: string,
+      newMessage?: MessageWithFiles | null,
+      runId?: string,
+    ) => {
       if (threadId) {
         queryClient.setQueryData(
           messagesWithFilesQuery(organization.id, project.id, threadId)
@@ -206,12 +216,30 @@ export function ChatProvider({
             const existingIndex = messages?.findIndex(
               (item) => item.id === newMessage.id,
             );
-            if (existingIndex) {
+            if (existingIndex && existingIndex !== -1) {
               return messages?.toSpliced(existingIndex, 1, newMessage);
             }
             return messages ? [...messages, newMessage] : [newMessage];
           },
         );
+        queryClient.invalidateQueries({
+          queryKey: messagesWithFilesQuery(
+            organization.id,
+            project.id,
+            threadId,
+          ).queryKey,
+        });
+        if (runId) {
+          queryClient.invalidateQueries({
+            queryKey: runStepsQuery(
+              organization.id,
+              project.id,
+              threadId,
+              runId,
+              PLAN_STEPS_QUERY_PARAMS,
+            ).queryKey,
+          });
+        }
       }
     },
     [project.id, organization.id, queryClient],
@@ -642,7 +670,11 @@ export function ChatProvider({
             },
           },
           onMessageCompleted: (response) => {
-            setMessagesWithFilesQueryData(thread?.id, response.data);
+            setMessagesWithFilesQueryData(
+              thread?.id,
+              response.data,
+              response.data?.run_id,
+            );
 
             if (files.length > 0) {
               queryClient.invalidateQueries({

diff --git a/src/modules/users/useUpdateUser.ts b/src/modules/users/useUpdateUser.ts
@@ -16,7 +16,7 @@
 
 import { updateUser } from '@/app/api/users';
 import { UserUpdateBody } from '@/app/api/users/types';
-import { encodeMetadata } from '@/app/api/utils';
+import { decodeMetadata } from '@/app/api/utils';
 import { UserMetadata } from '@/store/user-profile/types';
 import { useMutation } from '@tanstack/react-query';
 import { useSession } from 'next-auth/react';
@@ -31,7 +31,7 @@ export function useUpdateUser() {
         ...session,
         userProfile: {
           ...session?.userProfile,
-          metadata: encodeMetadata<UserMetadata>(result?.metadata),
+          metadata: decodeMetadata<UserMetadata>(result?.metadata),
         },
       });
     },