Fix MCP server start/stop for Tasks

src-tauri/src/mcp.rs

@@ -213,29 +213,30 @@ pub async fn rename_server(
     state: tauri::State<'_, State>,
 ) -> Result<()> {
     let mut sessions = state.sessions.lock().await;
-    
+
     if let Some(session) = sessions.get_mut(&session_id) {
         if let Some(server) = session.mcp_servers.get_mut(&old_name) {
             // Update the server's name
             server.set_name(new_name.clone());
-            
+
             // Update the tools mapping
-            let tools_to_update: Vec<String> = session.tools
+            let tools_to_update: Vec<String> = session
+                .tools
                 .iter()
                 .filter(|(_, server_name)| *server_name == &old_name)
                 .map(|(tool_name, _)| tool_name.clone())
                 .collect();
-                
+
             for tool_name in tools_to_update {
                 session.tools.insert(tool_name, new_name.clone());
             }
-            
+
             // Move the server to the new name in the map
             if let Some(server) = session.mcp_servers.remove(&old_name) {
                 session.mcp_servers.insert(new_name, server);
             }
         }
     }
-    
+
     Ok(())
 }

src/components/Task.svelte

@@ -51,6 +51,11 @@
     async function setModel(_model: Model) {
         task.engineId = Number(_model.engineId);
         task.model = String(_model.id);
+
+        if (isEdit) {
+            await onsave?.();
+        }
+
         model = _model;
     }
 
@@ -75,8 +80,8 @@
         <Toggle
             label={mcpServer.name}
             value={hasMcpServer(mcpServer) ? 'on' : 'off'}
-            onEnable={() => addMcpServer(mcpServer)}
-            onDisable={() => removeMcpServer(mcpServer)}
+            onEnable={async () => await addMcpServer(mcpServer)}
+            onDisable={async () => await removeMcpServer(mcpServer)}
         />
     </Flex>
 {/snippet}

src/lib/closables.ts

@@ -28,7 +28,7 @@ export default {
 
     close(e: Event) {
         closables.forEach(([ele, fn]) => {
-            if (!ele.contains?.(e.target as Node)) {
+            if (!ele?.contains?.(e.target as Node)) {
                 fn();
             }
         });

src/lib/models/session.svelte.ts

@@ -79,6 +79,20 @@ export default class Session extends Base<Row>('sessions') {
         return Message.where({ sessionId: this.id });
     }
 
+    get mcpServers(): McpServer[] {
+        return (this.config.enabledMcpServers || [])
+            .map(name => McpServer.findBy({ name }))
+            .compact();
+    }
+
+    async start() {
+        await Promise.all(this.mcpServers.map(async server => await server.start(this)));
+    }
+
+    async stop() {
+        await Promise.all(this.mcpServers.map(async server => await server.stop(this)));
+    }
+
     async tools(): Promise<Tool[]> {
         if (!this.id || !this.config?.model) {
             return [];

src/lib/tasks.ts

@@ -1,8 +1,9 @@
+import type { Mode } from 'highlight.js';
 import moment from 'moment';
 
 import { dispatch } from '$lib/dispatch';
 import { info } from '$lib/logger';
-import { Engine, Session, Task, TaskRun } from '$lib/models';
+import { Engine, Model, Session, Task, TaskRun } from '$lib/models';
 import { State } from '$lib/models/task-run.svelte';
 
 const TASK_APP_ID = 2;
@@ -62,21 +63,22 @@ export async function execute(task: Task): Promise<TaskRun | undefined> {
         ephemeral: true,
     });
 
-    if (!session) {
-        return;
-    }
-
-    task.mcpServers.forEach(async server => {
-        await session.addMcpServer(server);
-        await server.start(session);
-    });
-
     const run = await TaskRun.create({
         taskId: task.id,
         sessionId: session.id,
     });
 
-    dispatch(session, model, task.prompt)
+    // Run without awaiting, so we can return immediately, render the UI,
+    // and wait for result asynchronously.
+    _execute(run, model);
+
+    return run;
+}
+
+async function _execute(run: TaskRun, model: Model) {
+    await run.session.start();
+
+    await dispatch(run.session, model, run.task.prompt)
         .then(_ => {
             run.state = State.Success;
             run.save();
@@ -87,5 +89,5 @@ export async function execute(task: Task): Promise<TaskRun | undefined> {
             run.save();
         });
 
-    return run;
+    await run.session.stop();
 }

src/routes/tasks/[task_id]/edit/+page.svelte

@@ -7,7 +7,7 @@
     const taskId = Number(page.params.task_id);
     const task: Task = $derived(Task.find(taskId));
 
-    async function save(): Promise<Task> {
+    async function save() {
         return await task.save();
     }
 </script>