fix(mcp): ensure CLI shutdown cleanup
This commit is contained in:
parent
1b3d102684
commit
c1d3db1ad3
2 changed files with 65 additions and 2 deletions
|
|
@ -19,6 +19,14 @@ import (
|
|||
var workspaceFlag string
|
||||
var unrestrictedFlag bool
|
||||
|
||||
var newMCPService = mcp.New
|
||||
var runMCPService = func(svc *mcp.Service, ctx context.Context) error {
|
||||
return svc.Run(ctx)
|
||||
}
|
||||
var shutdownMCPService = func(svc *mcp.Service, ctx context.Context) error {
|
||||
return svc.Shutdown(ctx)
|
||||
}
|
||||
|
||||
var mcpCmd = &cli.Command{
|
||||
Use: "mcp",
|
||||
Short: "MCP server for AI tool integration",
|
||||
|
|
@ -87,10 +95,13 @@ func runServe() error {
|
|||
}
|
||||
|
||||
// Create the MCP service
|
||||
svc, err := mcp.New(opts)
|
||||
svc, err := newMCPService(opts)
|
||||
if err != nil {
|
||||
return cli.Wrap(err, "create MCP service")
|
||||
}
|
||||
defer func() {
|
||||
_ = shutdownMCPService(svc, context.Background())
|
||||
}()
|
||||
|
||||
// Set up signal handling for clean shutdown
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
|
@ -105,5 +116,5 @@ func runServe() error {
|
|||
}()
|
||||
|
||||
// Run the server (blocks until context cancelled or error)
|
||||
return svc.Run(ctx)
|
||||
return runMCPService(svc, ctx)
|
||||
}
|
||||
|
|
|
|||
52
cmd/mcpcmd/cmd_mcp_test.go
Normal file
52
cmd/mcpcmd/cmd_mcp_test.go
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
package mcpcmd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"dappco.re/go/mcp/pkg/mcp"
|
||||
)
|
||||
|
||||
func TestRunServe_Good_ShutsDownService(t *testing.T) {
|
||||
oldNew := newMCPService
|
||||
oldRun := runMCPService
|
||||
oldShutdown := shutdownMCPService
|
||||
oldWorkspace := workspaceFlag
|
||||
oldUnrestricted := unrestrictedFlag
|
||||
|
||||
t.Cleanup(func() {
|
||||
newMCPService = oldNew
|
||||
runMCPService = oldRun
|
||||
shutdownMCPService = oldShutdown
|
||||
workspaceFlag = oldWorkspace
|
||||
unrestrictedFlag = oldUnrestricted
|
||||
})
|
||||
|
||||
workspaceFlag = ""
|
||||
unrestrictedFlag = false
|
||||
|
||||
var runCalled bool
|
||||
var shutdownCalled bool
|
||||
|
||||
newMCPService = func(opts mcp.Options) (*mcp.Service, error) {
|
||||
return mcp.New(mcp.Options{})
|
||||
}
|
||||
runMCPService = func(svc *mcp.Service, ctx context.Context) error {
|
||||
runCalled = true
|
||||
return nil
|
||||
}
|
||||
shutdownMCPService = func(svc *mcp.Service, ctx context.Context) error {
|
||||
shutdownCalled = true
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := runServe(); err != nil {
|
||||
t.Fatalf("runServe() returned error: %v", err)
|
||||
}
|
||||
if !runCalled {
|
||||
t.Fatal("expected runMCPService to be called")
|
||||
}
|
||||
if !shutdownCalled {
|
||||
t.Fatal("expected shutdownMCPService to be called")
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue