feat: permission dialog for tool call requests

This commit is contained in:
Louis 2025-04-22 22:55:04 +07:00
parent 1a3cc64a7e
commit e1995a3ccb
No known key found for this signature in database
GPG Key ID: 44FA9F4D33C37DE2
8 changed files with 230 additions and 74 deletions

View File

@ -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 (
<Modal
title={<span>Allow tool from {modalProps.toolName} (local)?</span>}
open={isOpen}
onOpenChange={(open) => {
setIsOpen(!isOpen)
if(!open)
handleCancel()
}}
content={
<div>
<p className="text-[hsla(var(--text-secondary))]">
Malicious MCP servers or conversation content could potentially
trick Jan into attempting harmful actions through your installed
tools. Review each action carefully before approving.
</p>
<div className="mt-4 flex justify-end gap-x-2">
<ModalClose asChild>
<Button
theme="primary"
onClick={() => {
setApprovedToolsAtom((prev) => {
const newState = { ...prev }
if (!newState[modalProps.threadId]) {
newState[modalProps.threadId] = []
}
if (
!newState[modalProps.threadId].includes(
modalProps.toolName
)
) {
newState[modalProps.threadId].push(modalProps.toolName)
}
return newState
})
handleConfirm()
}}
autoFocus
>
Allow for this chat
</Button>
</ModalClose>
<ModalClose asChild>
<Button theme="primary" onClick={handleConfirm} autoFocus>
Allow once
</Button>
</ModalClose>
<ModalClose asChild onClick={handleCancel}>
<Button theme="ghost">Deny</Button>
</ModalClose>
</div>
</div>
}
/>
)
}, [isOpen, modalProps])
return { showModal, PromiseModal }
}

View File

@ -72,6 +72,11 @@ export const threadDataReadyAtom = atomWithStorage<boolean>(
*/ */
export const threadModelParamsAtom = atom<Record<string, ModelParams>>({}) export const threadModelParamsAtom = atom<Record<string, ModelParams>>({})
/**
* Store the tool call approval thread id
*/
export const approvedThreadToolsAtom = atom<Record<string, string[]>>({})
//// End Thread Atom //// End Thread Atom
/// Active Thread Atom /// Active Thread Atom

View File

@ -55,6 +55,7 @@ import {
import { selectedModelAtom } from '@/helpers/atoms/Model.atom' import { selectedModelAtom } from '@/helpers/atoms/Model.atom'
import { import {
activeThreadAtom, activeThreadAtom,
approvedThreadToolsAtom,
engineParamsUpdateAtom, engineParamsUpdateAtom,
getActiveThreadModelParamsAtom, getActiveThreadModelParamsAtom,
isGeneratingResponseAtom, isGeneratingResponseAtom,
@ -65,7 +66,9 @@ import { ModelTool } from '@/types/model'
export const reloadModelAtom = atom(false) export const reloadModelAtom = atom(false)
export default function useSendChatMessage() { export default function useSendChatMessage(
showModal?: (toolName: string, threadId: string) => Promise<unknown>
) {
const activeThread = useAtomValue(activeThreadAtom) const activeThread = useAtomValue(activeThreadAtom)
const activeAssistant = useAtomValue(activeAssistantAtom) const activeAssistant = useAtomValue(activeAssistantAtom)
const addNewMessage = useSetAtom(addNewMessageAtom) const addNewMessage = useSetAtom(addNewMessageAtom)
@ -74,6 +77,7 @@ export default function useSendChatMessage() {
const setCurrentPrompt = useSetAtom(currentPromptAtom) const setCurrentPrompt = useSetAtom(currentPromptAtom)
const deleteMessage = useSetAtom(deleteMessageAtom) const deleteMessage = useSetAtom(deleteMessageAtom)
const setEditPrompt = useSetAtom(editPromptAtom) const setEditPrompt = useSetAtom(editPromptAtom)
const approvedTools = useAtomValue(approvedThreadToolsAtom)
const currentMessages = useAtomValue(getCurrentChatMessagesAtom) const currentMessages = useAtomValue(getCurrentChatMessagesAtom)
const selectedModel = useAtomValue(selectedModelAtom) const selectedModel = useAtomValue(selectedModelAtom)
@ -480,10 +484,25 @@ export default function useSendChatMessage() {
} }
events.emit(MessageEvent.OnMessageUpdate, message) events.emit(MessageEvent.OnMessageUpdate, message)
const result = await window.core.api.callTool({ const approved =
toolName: toolCall.function.name, approvedTools[message.thread_id]?.includes(toolCall.function.name) ||
arguments: JSON.parse(toolCall.function.arguments), (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 if (result.error) break
message.metadata = { message.metadata = {

View File

@ -1,5 +1,5 @@
/* eslint-disable @typescript-eslint/no-explicit-any */ /* 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 { MessageStatus } from '@janhq/core'
import { useAtom, useAtomValue } from 'jotai' import { useAtom, useAtomValue } from 'jotai'
@ -28,6 +28,7 @@ import {
getActiveThreadIdAtom, getActiveThreadIdAtom,
activeSettingInputBoxAtom, activeSettingInputBoxAtom,
} from '@/helpers/atoms/Thread.atom' } from '@/helpers/atoms/Thread.atom'
import { ChatContext } from '../../ThreadCenterPanel'
type CustomElement = { type CustomElement = {
type: 'paragraph' | 'code' | null type: 'paragraph' | 'code' | null
@ -77,7 +78,8 @@ const RichTextEditor = ({
const activeThreadId = useAtomValue(getActiveThreadIdAtom) const activeThreadId = useAtomValue(getActiveThreadIdAtom)
const activeSettingInputBox = useAtomValue(activeSettingInputBoxAtom) const activeSettingInputBox = useAtomValue(activeSettingInputBoxAtom)
const messages = useAtomValue(getCurrentChatMessagesAtom) const messages = useAtomValue(getCurrentChatMessagesAtom)
const { sendChatMessage } = useSendChatMessage() const { showApprovalModal } = useContext(ChatContext)
const { sendChatMessage } = useSendChatMessage(showApprovalModal)
const { stopInference } = useActiveModel() const { stopInference } = useActiveModel()
const selectedModel = useAtomValue(selectedModelAtom) const selectedModel = useAtomValue(selectedModelAtom)
const largeContentThreshold = 1000 const largeContentThreshold = 1000

View File

@ -1,5 +1,5 @@
/* eslint-disable @typescript-eslint/no-explicit-any */ /* 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' import { InferenceEngine } from '@janhq/core'
@ -11,7 +11,6 @@ import {
badgeVariants, badgeVariants,
Badge, Badge,
Modal, Modal,
ModalClose,
} from '@janhq/joi' } from '@janhq/joi'
import { useAtom, useAtomValue } from 'jotai' import { useAtom, useAtomValue } from 'jotai'
import { import {
@ -56,6 +55,7 @@ import {
} from '@/helpers/atoms/Thread.atom' } from '@/helpers/atoms/Thread.atom'
import { activeTabThreadRightPanelAtom } from '@/helpers/atoms/ThreadRightPanel.atom' import { activeTabThreadRightPanelAtom } from '@/helpers/atoms/ThreadRightPanel.atom'
import { ModelTool } from '@/types/model' import { ModelTool } from '@/types/model'
import { ChatContext } from '../../ThreadCenterPanel'
const ChatInput = () => { const ChatInput = () => {
const activeThread = useAtomValue(activeThreadAtom) const activeThread = useAtomValue(activeThreadAtom)
@ -66,7 +66,8 @@ const ChatInput = () => {
const [activeSettingInputBox, setActiveSettingInputBox] = useAtom( const [activeSettingInputBox, setActiveSettingInputBox] = useAtom(
activeSettingInputBoxAtom activeSettingInputBoxAtom
) )
const { sendChatMessage } = useSendChatMessage() const { showApprovalModal } = useContext(ChatContext)
const { sendChatMessage } = useSendChatMessage(showApprovalModal)
const selectedModel = useAtomValue(selectedModelAtom) const selectedModel = useAtomValue(selectedModelAtom)
const activeThreadId = useAtomValue(getActiveThreadIdAtom) const activeThreadId = useAtomValue(getActiveThreadIdAtom)

View File

@ -1,5 +1,5 @@
/* eslint-disable @typescript-eslint/no-explicit-any */ /* eslint-disable @typescript-eslint/no-explicit-any */
import { useEffect, useRef, useState } from 'react' import { useContext, useEffect, useRef, useState } from 'react'
import { import {
ConversationalExtension, ConversationalExtension,
@ -31,6 +31,7 @@ import {
activeThreadAtom, activeThreadAtom,
getActiveThreadIdAtom, getActiveThreadIdAtom,
} from '@/helpers/atoms/Thread.atom' } from '@/helpers/atoms/Thread.atom'
import { ChatContext } from '../../ThreadCenterPanel'
type Props = { type Props = {
message: ThreadMessage message: ThreadMessage
@ -42,7 +43,8 @@ const EditChatInput: React.FC<Props> = ({ message }) => {
const messages = useAtomValue(getCurrentChatMessagesAtom) const messages = useAtomValue(getCurrentChatMessagesAtom)
const [editPrompt, setEditPrompt] = useAtom(editPromptAtom) const [editPrompt, setEditPrompt] = useAtom(editPromptAtom)
const { sendChatMessage } = useSendChatMessage() const { showApprovalModal } = useContext(ChatContext)
const { sendChatMessage } = useSendChatMessage(showApprovalModal)
const setMessages = useSetAtom(setConvoMessagesAtom) const setMessages = useSetAtom(setConvoMessagesAtom)
const activeThreadId = useAtomValue(getActiveThreadIdAtom) const activeThreadId = useAtomValue(getActiveThreadIdAtom)
const spellCheck = useAtomValue(spellCheckAtom) const spellCheck = useAtomValue(spellCheckAtom)

View File

@ -1,4 +1,4 @@
import { useCallback } from 'react' import { useCallback, useContext } from 'react'
import { import {
MessageStatus, MessageStatus,
@ -34,13 +34,15 @@ import {
updateThreadAtom, updateThreadAtom,
updateThreadStateLastMessageAtom, updateThreadStateLastMessageAtom,
} from '@/helpers/atoms/Thread.atom' } from '@/helpers/atoms/Thread.atom'
import { ChatContext } from '../../ThreadCenterPanel'
const MessageToolbar = ({ message }: { message: ThreadMessage }) => { const MessageToolbar = ({ message }: { message: ThreadMessage }) => {
const deleteMessage = useSetAtom(deleteMessageAtom) const deleteMessage = useSetAtom(deleteMessageAtom)
const setEditMessage = useSetAtom(editMessageAtom) const setEditMessage = useSetAtom(editMessageAtom)
const thread = useAtomValue(activeThreadAtom) const thread = useAtomValue(activeThreadAtom)
const messages = useAtomValue(getCurrentChatMessagesAtom) const messages = useAtomValue(getCurrentChatMessagesAtom)
const { resendChatMessage } = useSendChatMessage() const { showApprovalModal } = useContext(ChatContext)
const { resendChatMessage } = useSendChatMessage(showApprovalModal)
const clipboard = useClipboard({ timeout: 1000 }) const clipboard = useClipboard({ timeout: 1000 })
const updateThreadLastMessage = useSetAtom(updateThreadStateLastMessageAtom) const updateThreadLastMessage = useSetAtom(updateThreadStateLastMessageAtom)
const updateThread = useSetAtom(updateThreadAtom) const updateThread = useSetAtom(updateThreadAtom)

View File

@ -1,6 +1,6 @@
/* eslint-disable @typescript-eslint/naming-convention */ /* 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' import { Accept, useDropzone } from 'react-dropzone'
@ -36,6 +36,7 @@ import {
engineParamsUpdateAtom, engineParamsUpdateAtom,
isGeneratingResponseAtom, isGeneratingResponseAtom,
} from '@/helpers/atoms/Thread.atom' } from '@/helpers/atoms/Thread.atom'
import { useTollCallPromiseModal } from '@/containers/ToolCallApprovalModal'
const renderError = (code: string) => { const renderError = (code: string) => {
switch (code) { switch (code) {
@ -53,6 +54,16 @@ const renderError = (code: string) => {
} }
} }
interface ChatContextType {
showApprovalModal:
| ((toolName: string, threadId: string) => Promise<unknown>)
| undefined
}
export const ChatContext = createContext<ChatContextType>({
showApprovalModal: undefined,
})
const ThreadCenterPanel = () => { const ThreadCenterPanel = () => {
const [dragRejected, setDragRejected] = useState({ code: '' }) const [dragRejected, setDragRejected] = useState({ code: '' })
const setFileUpload = useSetAtom(fileUploadAtom) const setFileUpload = useSetAtom(fileUploadAtom)
@ -71,6 +82,7 @@ const ThreadCenterPanel = () => {
: { : {
'application/pdf': ['.pdf'], 'application/pdf': ['.pdf'],
} }
const { showModal, PromiseModal } = useTollCallPromiseModal()
const { getRootProps, isDragReject } = useDropzone({ const { getRootProps, isDragReject } = useDropzone({
noClick: true, noClick: true,
@ -158,73 +170,82 @@ const ThreadCenterPanel = () => {
const isGeneratingResponse = useAtomValue(isGeneratingResponseAtom) const isGeneratingResponse = useAtomValue(isGeneratingResponseAtom)
return ( return (
<CenterPanelContainer> <ChatContext.Provider
<div value={{
className="relative flex h-full w-full flex-col outline-none" showApprovalModal: showModal,
{...getRootProps()} }}
> >
{dragOver && ( <CenterPanelContainer>
<div className="absolute z-50 mx-auto h-full w-full p-8 backdrop-blur-lg"> <div
<div className="relative flex h-full w-full flex-col outline-none"
className={twMerge( {...getRootProps()}
'flex h-full w-full items-center justify-center rounded-lg border border-dashed border-[hsla(var(--primary-bg))]', >
isDragReject && 'border-[hsla(var(--destructive-bg))]' {dragOver && (
)} <div className="absolute z-50 mx-auto h-full w-full p-8 backdrop-blur-lg">
> <div
<div className="mx-auto w-1/2 text-center"> className={twMerge(
<div className="mx-auto inline-flex h-12 w-12 items-center justify-center rounded-full"> 'flex h-full w-full items-center justify-center rounded-lg border border-dashed border-[hsla(var(--primary-bg))]',
<UploadCloudIcon isDragReject && 'border-[hsla(var(--destructive-bg))]'
size={24} )}
className="text-[hsla(var(--primary-bg))]" >
/> <div className="mx-auto w-1/2 text-center">
</div> <div className="mx-auto inline-flex h-12 w-12 items-center justify-center rounded-full">
<div className="mt-4 text-[hsla(var(--primary-bg))]"> <UploadCloudIcon
<h6 className="font-bold"> size={24}
{isDragReject className="text-[hsla(var(--primary-bg))]"
? `Currently, we only support 1 attachment at the same time with ${ />
activeAssistant?.model.settings?.mmproj </div>
? 'PDF, JPEG, JPG, PNG' <div className="mt-4 text-[hsla(var(--primary-bg))]">
: 'PDF' <h6 className="font-bold">
} format` {isDragReject
: 'Drop file here'} ? `Currently, we only support 1 attachment at the same time with ${
</h6> activeAssistant?.model.settings?.mmproj
{!isDragReject && ( ? 'PDF, JPEG, JPG, PNG'
<p className="mt-2"> : 'PDF'
{activeAssistant?.model.settings?.mmproj } format`
? 'PDF, JPEG, JPG, PNG' : 'Drop file here'}
: 'PDF'} </h6>
</p> {!isDragReject && (
)} <p className="mt-2">
{activeAssistant?.model.settings?.mmproj
? 'PDF, JPEG, JPG, PNG'
: 'PDF'}
</p>
)}
</div>
</div> </div>
</div> </div>
</div> </div>
</div>
)}
<div className={twMerge('flex h-full w-full flex-col justify-between')}>
{activeThread ? (
<div className="flex h-full w-full overflow-x-hidden">
<ChatBody />
</div>
) : (
<RequestDownloadModel />
)} )}
{!engineParamsUpdate && <ModelStart />}
{reloadModel && <ModelReload />}
{activeModel && isGeneratingResponse && <GenerateResponse />}
<div <div
className={twMerge( className={twMerge('flex h-full w-full flex-col justify-between')}
'mx-auto w-full',
chatWidth === 'compact' && 'max-w-[700px]'
)}
> >
<ChatInput /> {activeThread ? (
<div className="flex h-full w-full overflow-x-hidden">
<ChatBody />
</div>
) : (
<RequestDownloadModel />
)}
{!engineParamsUpdate && <ModelStart />}
{reloadModel && <ModelReload />}
{activeModel && isGeneratingResponse && <GenerateResponse />}
<div
className={twMerge(
'mx-auto w-full',
chatWidth === 'compact' && 'max-w-[700px]'
)}
>
<ChatInput />
</div>
</div> </div>
</div> </div>
</div> </CenterPanelContainer>
</CenterPanelContainer> <PromiseModal />
</ChatContext.Provider>
) )
} }