From e1995a3ccb54ed9819f6183dcbc7fb2d4ae21f58 Mon Sep 17 00:00:00 2001 From: Louis Date: Tue, 22 Apr 2025 22:55:04 +0700 Subject: [PATCH] feat: permission dialog for tool call requests --- .../ToolCallApprovalModal/index.tsx | 104 +++++++++++++ web/helpers/atoms/Thread.atom.ts | 5 + web/hooks/useSendChatMessage.ts | 29 +++- .../ChatInput/RichTextEditor.tsx | 6 +- .../ThreadCenterPanel/ChatInput/index.tsx | 7 +- .../ThreadCenterPanel/EditChatInput/index.tsx | 6 +- .../MessageToolbar/index.tsx | 6 +- .../Thread/ThreadCenterPanel/index.tsx | 141 ++++++++++-------- 8 files changed, 230 insertions(+), 74 deletions(-) create mode 100644 web/containers/ToolCallApprovalModal/index.tsx diff --git a/web/containers/ToolCallApprovalModal/index.tsx b/web/containers/ToolCallApprovalModal/index.tsx new file mode 100644 index 000000000..2097278ad --- /dev/null +++ b/web/containers/ToolCallApprovalModal/index.tsx @@ -0,0 +1,104 @@ +import { memo, useCallback, useState } from 'react' + +import { Button, Modal, ModalClose } from '@janhq/joi' +import { useSetAtom } from 'jotai' +import { approvedThreadToolsAtom } from '@/helpers/atoms/Thread.atom' + +export function useTollCallPromiseModal() { + const [isOpen, setIsOpen] = useState(false) + const setApprovedToolsAtom = useSetAtom(approvedThreadToolsAtom) + const [modalProps, setModalProps] = useState<{ + toolName: string + threadId: string + resolveRef: ((value: unknown) => void) | null + }>({ + toolName: '', + threadId: '', + resolveRef: null, + }) + + // Function to open the modal and return a Promise + const showModal = useCallback((toolName: string, threadId: string) => { + return new Promise((resolve) => { + setModalProps({ + toolName, + threadId, + resolveRef: resolve, + }) + setIsOpen(true) + }) + }, []) + + const PromiseModal = useCallback(() => { + const handleConfirm = () => { + setIsOpen(false) + if (modalProps.resolveRef) { + modalProps.resolveRef(true) + } + } + + const handleCancel = () => { + setIsOpen(false) + if (modalProps.resolveRef) { + modalProps.resolveRef(false) + } + } + + return ( + Allow tool from {modalProps.toolName} (local)?} + open={isOpen} + onOpenChange={(open) => { + setIsOpen(!isOpen) + if(!open) + handleCancel() + }} + content={ +
+

+ Malicious MCP servers or conversation content could potentially + trick Jan into attempting harmful actions through your installed + tools. Review each action carefully before approving. +

+
+ + + + + + + + + +
+
+ } + /> + ) + }, [isOpen, modalProps]) + return { showModal, PromiseModal } +} diff --git a/web/helpers/atoms/Thread.atom.ts b/web/helpers/atoms/Thread.atom.ts index 474dadeba..d60e9dded 100644 --- a/web/helpers/atoms/Thread.atom.ts +++ b/web/helpers/atoms/Thread.atom.ts @@ -72,6 +72,11 @@ export const threadDataReadyAtom = atomWithStorage( */ export const threadModelParamsAtom = atom>({}) +/** + * Store the tool call approval thread id + */ +export const approvedThreadToolsAtom = atom>({}) + //// End Thread Atom /// Active Thread Atom diff --git a/web/hooks/useSendChatMessage.ts b/web/hooks/useSendChatMessage.ts index 49e0d3e5b..7a25ba797 100644 --- a/web/hooks/useSendChatMessage.ts +++ b/web/hooks/useSendChatMessage.ts @@ -55,6 +55,7 @@ import { import { selectedModelAtom } from '@/helpers/atoms/Model.atom' import { activeThreadAtom, + approvedThreadToolsAtom, engineParamsUpdateAtom, getActiveThreadModelParamsAtom, isGeneratingResponseAtom, @@ -65,7 +66,9 @@ import { ModelTool } from '@/types/model' export const reloadModelAtom = atom(false) -export default function useSendChatMessage() { +export default function useSendChatMessage( + showModal?: (toolName: string, threadId: string) => Promise +) { const activeThread = useAtomValue(activeThreadAtom) const activeAssistant = useAtomValue(activeAssistantAtom) const addNewMessage = useSetAtom(addNewMessageAtom) @@ -74,6 +77,7 @@ export default function useSendChatMessage() { const setCurrentPrompt = useSetAtom(currentPromptAtom) const deleteMessage = useSetAtom(deleteMessageAtom) const setEditPrompt = useSetAtom(editPromptAtom) + const approvedTools = useAtomValue(approvedThreadToolsAtom) const currentMessages = useAtomValue(getCurrentChatMessagesAtom) const selectedModel = useAtomValue(selectedModelAtom) @@ -480,10 +484,25 @@ export default function useSendChatMessage() { } events.emit(MessageEvent.OnMessageUpdate, message) - const result = await window.core.api.callTool({ - toolName: toolCall.function.name, - arguments: JSON.parse(toolCall.function.arguments), - }) + const approved = + approvedTools[message.thread_id]?.includes(toolCall.function.name) || + (showModal + ? await showModal(toolCall.function.name, message.thread_id) + : true) + + const result = approved + ? await window.core.api.callTool({ + toolName: toolCall.function.name, + arguments: JSON.parse(toolCall.function.arguments), + }) + : { + content: [ + { + type: 'text', + text: 'The user has chosen to disallow the tool call.', + }, + ], + } if (result.error) break message.metadata = { diff --git a/web/screens/Thread/ThreadCenterPanel/ChatInput/RichTextEditor.tsx b/web/screens/Thread/ThreadCenterPanel/ChatInput/RichTextEditor.tsx index e6ee94c9a..4bf43c381 100644 --- a/web/screens/Thread/ThreadCenterPanel/ChatInput/RichTextEditor.tsx +++ b/web/screens/Thread/ThreadCenterPanel/ChatInput/RichTextEditor.tsx @@ -1,5 +1,5 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ -import { useCallback, useEffect, useMemo, useRef, ClipboardEvent } from 'react' +import { useCallback, useEffect, useMemo, useRef, ClipboardEvent, useContext } from 'react' import { MessageStatus } from '@janhq/core' import { useAtom, useAtomValue } from 'jotai' @@ -28,6 +28,7 @@ import { getActiveThreadIdAtom, activeSettingInputBoxAtom, } from '@/helpers/atoms/Thread.atom' +import { ChatContext } from '../../ThreadCenterPanel' type CustomElement = { type: 'paragraph' | 'code' | null @@ -77,7 +78,8 @@ const RichTextEditor = ({ const activeThreadId = useAtomValue(getActiveThreadIdAtom) const activeSettingInputBox = useAtomValue(activeSettingInputBoxAtom) const messages = useAtomValue(getCurrentChatMessagesAtom) - const { sendChatMessage } = useSendChatMessage() + const { showApprovalModal } = useContext(ChatContext) + const { sendChatMessage } = useSendChatMessage(showApprovalModal) const { stopInference } = useActiveModel() const selectedModel = useAtomValue(selectedModelAtom) const largeContentThreshold = 1000 diff --git a/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx index 3e8c3f55c..b276c1b90 100644 --- a/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx +++ b/web/screens/Thread/ThreadCenterPanel/ChatInput/index.tsx @@ -1,5 +1,5 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ -import { useEffect, useRef, useState } from 'react' +import { useContext, useEffect, useRef, useState } from 'react' import { InferenceEngine } from '@janhq/core' @@ -11,7 +11,6 @@ import { badgeVariants, Badge, Modal, - ModalClose, } from '@janhq/joi' import { useAtom, useAtomValue } from 'jotai' import { @@ -56,6 +55,7 @@ import { } from '@/helpers/atoms/Thread.atom' import { activeTabThreadRightPanelAtom } from '@/helpers/atoms/ThreadRightPanel.atom' import { ModelTool } from '@/types/model' +import { ChatContext } from '../../ThreadCenterPanel' const ChatInput = () => { const activeThread = useAtomValue(activeThreadAtom) @@ -66,7 +66,8 @@ const ChatInput = () => { const [activeSettingInputBox, setActiveSettingInputBox] = useAtom( activeSettingInputBoxAtom ) - const { sendChatMessage } = useSendChatMessage() + const { showApprovalModal } = useContext(ChatContext) + const { sendChatMessage } = useSendChatMessage(showApprovalModal) const selectedModel = useAtomValue(selectedModelAtom) const activeThreadId = useAtomValue(getActiveThreadIdAtom) diff --git a/web/screens/Thread/ThreadCenterPanel/EditChatInput/index.tsx b/web/screens/Thread/ThreadCenterPanel/EditChatInput/index.tsx index 88f6a72b0..178551742 100644 --- a/web/screens/Thread/ThreadCenterPanel/EditChatInput/index.tsx +++ b/web/screens/Thread/ThreadCenterPanel/EditChatInput/index.tsx @@ -1,5 +1,5 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ -import { useEffect, useRef, useState } from 'react' +import { useContext, useEffect, useRef, useState } from 'react' import { ConversationalExtension, @@ -31,6 +31,7 @@ import { activeThreadAtom, getActiveThreadIdAtom, } from '@/helpers/atoms/Thread.atom' +import { ChatContext } from '../../ThreadCenterPanel' type Props = { message: ThreadMessage @@ -42,7 +43,8 @@ const EditChatInput: React.FC = ({ message }) => { const messages = useAtomValue(getCurrentChatMessagesAtom) const [editPrompt, setEditPrompt] = useAtom(editPromptAtom) - const { sendChatMessage } = useSendChatMessage() + const { showApprovalModal } = useContext(ChatContext) + const { sendChatMessage } = useSendChatMessage(showApprovalModal) const setMessages = useSetAtom(setConvoMessagesAtom) const activeThreadId = useAtomValue(getActiveThreadIdAtom) const spellCheck = useAtomValue(spellCheckAtom) diff --git a/web/screens/Thread/ThreadCenterPanel/MessageToolbar/index.tsx b/web/screens/Thread/ThreadCenterPanel/MessageToolbar/index.tsx index 47d2cfc63..0e81996ea 100644 --- a/web/screens/Thread/ThreadCenterPanel/MessageToolbar/index.tsx +++ b/web/screens/Thread/ThreadCenterPanel/MessageToolbar/index.tsx @@ -1,4 +1,4 @@ -import { useCallback } from 'react' +import { useCallback, useContext } from 'react' import { MessageStatus, @@ -34,13 +34,15 @@ import { updateThreadAtom, updateThreadStateLastMessageAtom, } from '@/helpers/atoms/Thread.atom' +import { ChatContext } from '../../ThreadCenterPanel' const MessageToolbar = ({ message }: { message: ThreadMessage }) => { const deleteMessage = useSetAtom(deleteMessageAtom) const setEditMessage = useSetAtom(editMessageAtom) const thread = useAtomValue(activeThreadAtom) const messages = useAtomValue(getCurrentChatMessagesAtom) - const { resendChatMessage } = useSendChatMessage() + const { showApprovalModal } = useContext(ChatContext) + const { resendChatMessage } = useSendChatMessage(showApprovalModal) const clipboard = useClipboard({ timeout: 1000 }) const updateThreadLastMessage = useSetAtom(updateThreadStateLastMessageAtom) const updateThread = useSetAtom(updateThreadAtom) diff --git a/web/screens/Thread/ThreadCenterPanel/index.tsx b/web/screens/Thread/ThreadCenterPanel/index.tsx index 02fd7c7fd..d2ee3a4e6 100644 --- a/web/screens/Thread/ThreadCenterPanel/index.tsx +++ b/web/screens/Thread/ThreadCenterPanel/index.tsx @@ -1,6 +1,6 @@ /* eslint-disable @typescript-eslint/naming-convention */ -import { memo, useEffect, useState } from 'react' +import { createContext, memo, useEffect, useState } from 'react' import { Accept, useDropzone } from 'react-dropzone' @@ -36,6 +36,7 @@ import { engineParamsUpdateAtom, isGeneratingResponseAtom, } from '@/helpers/atoms/Thread.atom' +import { useTollCallPromiseModal } from '@/containers/ToolCallApprovalModal' const renderError = (code: string) => { switch (code) { @@ -53,6 +54,16 @@ const renderError = (code: string) => { } } +interface ChatContextType { + showApprovalModal: + | ((toolName: string, threadId: string) => Promise) + | undefined +} + +export const ChatContext = createContext({ + showApprovalModal: undefined, +}) + const ThreadCenterPanel = () => { const [dragRejected, setDragRejected] = useState({ code: '' }) const setFileUpload = useSetAtom(fileUploadAtom) @@ -71,6 +82,7 @@ const ThreadCenterPanel = () => { : { 'application/pdf': ['.pdf'], } + const { showModal, PromiseModal } = useTollCallPromiseModal() const { getRootProps, isDragReject } = useDropzone({ noClick: true, @@ -158,73 +170,82 @@ const ThreadCenterPanel = () => { const isGeneratingResponse = useAtomValue(isGeneratingResponseAtom) return ( - -
- {dragOver && ( -
-
-
-
- -
-
-
- {isDragReject - ? `Currently, we only support 1 attachment at the same time with ${ - activeAssistant?.model.settings?.mmproj - ? 'PDF, JPEG, JPG, PNG' - : 'PDF' - } format` - : 'Drop file here'} -
- {!isDragReject && ( -

- {activeAssistant?.model.settings?.mmproj - ? 'PDF, JPEG, JPG, PNG' - : 'PDF'} -

- )} + + +
+ {dragOver && ( +
+
+
+
+ +
+
+
+ {isDragReject + ? `Currently, we only support 1 attachment at the same time with ${ + activeAssistant?.model.settings?.mmproj + ? 'PDF, JPEG, JPG, PNG' + : 'PDF' + } format` + : 'Drop file here'} +
+ {!isDragReject && ( +

+ {activeAssistant?.model.settings?.mmproj + ? 'PDF, JPEG, JPG, PNG' + : 'PDF'} +

+ )} +
-
- )} -
- {activeThread ? ( -
- -
- ) : ( - )} - - {!engineParamsUpdate && } - - {reloadModel && } - - {activeModel && isGeneratingResponse && }
- + {activeThread ? ( +
+ +
+ ) : ( + + )} + + {!engineParamsUpdate && } + + {reloadModel && } + + {activeModel && isGeneratingResponse && } +
+ +
-
- + + + ) }