From e09cd6b6d79f4c0b610f332fef53b55499d0a8bb Mon Sep 17 00:00:00 2001
From: suyao <sy20010504@gmail.com>
Date: Mon, 22 Sep 2025 15:05:03 +0800
Subject: [PATCH] Fix tool result handling and session creation flow

- Populate toolName in tool-result chunks from contentBlockState
- Add onSessionCreated callback to SessionModal for post-creation actions
- Return created session from useSessions hook and update SWR cache optimistically
---
 .../agents/services/claudecode/transform.ts   |  2 +-
 .../components/Popups/agent/SessionModal.tsx  | 74 +++++++++++--------
 src/renderer/src/hooks/agents/useSessions.ts  |  4 +-
 .../pages/home/Tabs/components/Sessions.tsx   |  1 +
 4 files changed, 49 insertions(+), 32 deletions(-)

diff --git a/src/main/services/agents/services/claudecode/transform.ts b/src/main/services/agents/services/claudecode/transform.ts
index 33d6e72226..34a3de17e0 100644
--- a/src/main/services/agents/services/claudecode/transform.ts
+++ b/src/main/services/agents/services/claudecode/transform.ts
@@ -118,7 +118,7 @@ function handleUserOrAssistantMessage(message: Extract<SDKMessage, { type: 'assi
           chunks.push({
             type: 'tool-result',
             toolCallId: block.tool_use_id,
-            toolName: '',
+            toolName: contentBlockState[block.tool_use_id].toolName,
             input: '',
             output: block.content
           })
diff --git a/src/renderer/src/components/Popups/agent/SessionModal.tsx b/src/renderer/src/components/Popups/agent/SessionModal.tsx
index 6e64e4a1aa..cf5750c124 100644
--- a/src/renderer/src/components/Popups/agent/SessionModal.tsx
+++ b/src/renderer/src/components/Popups/agent/SessionModal.tsx
@@ -46,6 +46,7 @@ const buildSessionForm = (existing?: AgentSessionEntity, agent?: AgentEntity): B
 interface BaseProps {
   agentId: string
   session?: AgentSessionEntity
+  onSessionCreated?: (session: AgentSessionEntity) => void
 }
 
 interface TriggerProps extends BaseProps {
@@ -73,7 +74,14 @@ type Props = TriggerProps | StateProps
  * @param onClose - Optional callback when modal closes. From useDisclosure.
  * @returns Modal component for agent creation/editing
  */
-export const SessionModal: React.FC<Props> = ({ agentId, session, trigger, isOpen: _isOpen, onClose: _onClose }) => {
+export const SessionModal: React.FC<Props> = ({
+  agentId,
+  session,
+  trigger,
+  isOpen: _isOpen,
+  onClose: _onClose,
+  onSessionCreated
+}) => {
   const { isOpen, onClose, onOpen } = useDisclosure({ isOpen: _isOpen, onClose: _onClose })
   const { t } = useTranslation()
   const loadingRef = useRef(false)
@@ -161,38 +169,43 @@ export const SessionModal: React.FC<Props> = ({ agentId, session, trigger, isOpe
         return
       }
 
-      if (isEditing(session)) {
-        if (!session) {
-          throw new Error('Agent is required for editing mode')
+      try {
+        if (isEditing(session)) {
+          if (!session) {
+            throw new Error('Agent is required for editing mode')
+          }
+
+          const updatePayload = {
+            id: session.id,
+            name: form.name,
+            description: form.description,
+            instructions: form.instructions,
+            model: form.model,
+            accessible_paths: [...form.accessible_paths]
+          } satisfies UpdateSessionForm
+
+          updateSession(updatePayload)
+          logger.debug('Updated agent', updatePayload)
+        } else {
+          const newSession = {
+            name: form.name,
+            description: form.description,
+            instructions: form.instructions,
+            model: form.model,
+            accessible_paths: [...form.accessible_paths]
+          } satisfies CreateSessionForm
+          const createdSession = await createSession(newSession)
+          if (createdSession) {
+            onSessionCreated?.(createdSession)
+          }
+          logger.debug('Added agent', newSession)
         }
 
-        const updatePayload = {
-          id: session.id,
-          name: form.name,
-          description: form.description,
-          instructions: form.instructions,
-          model: form.model,
-          accessible_paths: [...form.accessible_paths]
-        } satisfies UpdateSessionForm
-
-        updateSession(updatePayload)
-        logger.debug('Updated agent', updatePayload)
-      } else {
-        const newSession = {
-          name: form.name,
-          description: form.description,
-          instructions: form.instructions,
-          model: form.model,
-          accessible_paths: [...form.accessible_paths]
-        } satisfies CreateSessionForm
-        createSession(newSession)
-        logger.debug('Added agent', newSession)
+        // setTimeoutTimer('onCreateAgent', () => EventEmitter.emit(EVENT_NAMES.SHOW_ASSISTANTS), 0)
+        onClose()
+      } finally {
+        loadingRef.current = false
       }
-
-      loadingRef.current = false
-
-      // setTimeoutTimer('onCreateAgent', () => EventEmitter.emit(EVENT_NAMES.SHOW_ASSISTANTS), 0)
-      onClose()
     },
     [
       form.model,
@@ -202,6 +215,7 @@ export const SessionModal: React.FC<Props> = ({ agentId, session, trigger, isOpe
       form.accessible_paths,
       session,
       onClose,
+      onSessionCreated,
       t,
       updateSession,
       createSession
diff --git a/src/renderer/src/hooks/agents/useSessions.ts b/src/renderer/src/hooks/agents/useSessions.ts
index 1ed40d2f1b..9e83ac5afc 100644
--- a/src/renderer/src/hooks/agents/useSessions.ts
+++ b/src/renderer/src/hooks/agents/useSessions.ts
@@ -21,9 +21,11 @@ export const useSessions = (agentId: string) => {
     async (form: CreateSessionForm) => {
       try {
         const result = await client.createSession(agentId, form)
-        mutate((prev) => [...(prev ?? []), result])
+        await mutate((prev) => [...(prev ?? []), result], { revalidate: false })
+        return result
       } catch (error) {
         window.toast.error(formatErrorMessageWithPrefix(error, t('agent.session.create.error.failed')))
+        return undefined
       }
     },
     [agentId, client, mutate, t]
diff --git a/src/renderer/src/pages/home/Tabs/components/Sessions.tsx b/src/renderer/src/pages/home/Tabs/components/Sessions.tsx
index 2964e761e6..a8ccab7d00 100644
--- a/src/renderer/src/pages/home/Tabs/components/Sessions.tsx
+++ b/src/renderer/src/pages/home/Tabs/components/Sessions.tsx
@@ -66,6 +66,7 @@ const Sessions: React.FC<SessionsProps> = ({ agentId }) => {
         transition={{ duration: 0.2, delay: 0.1 }}>
         <SessionModal
           agentId={agentId}
+          onSessionCreated={(created) => setActiveSessionId(agentId, created.id)}
           trigger={{
             content: (
               <Button