From ca482b3edf3f74da0e278e844ee26d3cb8fe2ad3 Mon Sep 17 00:00:00 2001
From: Donnie Adams <donnie@acorn.io>
Date: Thu, 8 May 2025 20:49:13 -0400
Subject: [PATCH] enhance: add the ability to close an individual MCP server

Signed-off-by: Donnie Adams <donnie@acorn.io>
---
 pkg/engine/engine.go |  3 +++
 pkg/mcp/loader.go    | 26 ++++++++++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/pkg/engine/engine.go b/pkg/engine/engine.go
index c7867512..e3ff930a 100644
--- a/pkg/engine/engine.go
+++ b/pkg/engine/engine.go
@@ -350,6 +350,9 @@ func (e *Engine) Start(ctx Context, input string) (ret *Return, err error) {
 		}
 		select {
 		case <-ctx.userCancel:
+			if ret == nil {
+				ret = new(Return)
+			}
 			if ret.Result == nil {
 				ret.Result = new(string)
 			}
diff --git a/pkg/mcp/loader.go b/pkg/mcp/loader.go
index 72c0ebf4..86e2f8d2 100644
--- a/pkg/mcp/loader.go
+++ b/pkg/mcp/loader.go
@@ -125,6 +125,32 @@ func (l *Local) LoadTools(ctx context.Context, server ServerConfig, toolName str
 	return l.sessionToTools(ctx, session, toolName, allowedTools)
 }
 
+func (l *Local) ShutdownServer(server ServerConfig) error {
+	if l == nil {
+		return nil
+	}
+
+	id := hash.Digest(server)
+
+	l.lock.Lock()
+
+	if l.sessionCtx == nil {
+		l.lock.Unlock()
+		return nil
+	}
+
+	session := l.sessions[id]
+	delete(l.sessions, id)
+
+	l.lock.Unlock()
+
+	if session == nil {
+		return nil
+	}
+
+	return session.Client.Close()
+}
+
 func (l *Local) Close() error {
 	if l == nil {
 		return nil