From b5816bea5c4c1eec852fd9a66c15e3b23f89e769 Mon Sep 17 00:00:00 2001 From: rahulguptajss Date: Wed, 22 Apr 2026 12:55:56 +0530 Subject: [PATCH 1/5] ci: add session handling to agent --- integration/test/tools_test.go | 41 +++++++++++++++++++++++++++++----- 1 file changed, 35 insertions(+), 6 deletions(-) diff --git a/integration/test/tools_test.go b/integration/test/tools_test.go index 89abd32..f125305 100644 --- a/integration/test/tools_test.go +++ b/integration/test/tools_test.go @@ -7,18 +7,19 @@ import ( "encoding/json" "errors" "fmt" - "github.com/carlmjohnson/requests" - "github.com/modelcontextprotocol/go-sdk/mcp" - "github.com/netapp/ontap-mcp/config" - "github.com/netapp/ontap-mcp/ontap" - "github.com/openai/openai-go/v3" - "github.com/openai/openai-go/v3/option" "log/slog" "net/http" "os" "strings" "testing" "time" + + "github.com/carlmjohnson/requests" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/netapp/ontap-mcp/config" + "github.com/netapp/ontap-mcp/ontap" + "github.com/openai/openai-go/v3" + "github.com/openai/openai-go/v3/option" ) const ( @@ -79,6 +80,7 @@ type Agent struct { openaiClient openai.Client mcpSession *mcp.ClientSession mcpClient *mcp.Client + mcpTransport *mcp.StreamableClientTransport tools []*mcp.Tool model string } @@ -118,6 +120,7 @@ func NewAgent(llmUserName, llmToken, llmBaseURL, openaiModel, mcpServerURL strin openaiClient: openaiClient, mcpSession: session, mcpClient: mcpClient, + mcpTransport: mcpTransport, tools: toolsResult.Tools, model: openaiModel, }, nil @@ -135,11 +138,32 @@ func (a *Agent) convertMCPToolsToOpenAI() []openai.ChatCompletionToolUnionParam return tools } +func (a *Agent) reconnect(ctx context.Context) error { + _ = a.mcpSession.Close() + session, err := a.mcpClient.Connect(ctx, a.mcpTransport, nil) + if err != nil { + return fmt.Errorf("failed to reconnect to MCP server: %w", err) + } + a.mcpSession = session + slog.Info("Reconnected to MCP server") + return nil +} + func (a *Agent) callMCPTool(ctx context.Context, toolName string, arguments map[string]any) (string, error) { result, err := a.mcpSession.CallTool(ctx, &mcp.CallToolParams{ Name: toolName, Arguments: arguments, }) + if err != nil && strings.Contains(err.Error(), "connection closed") { + slog.Warn("MCP session dropped, reconnecting", slog.Any("error", err)) + if reconnErr := a.reconnect(ctx); reconnErr != nil { + return "", fmt.Errorf("failed to call tool: %w (reconnect failed: %w)", err, reconnErr) + } + result, err = a.mcpSession.CallTool(ctx, &mcp.CallToolParams{ + Name: toolName, + Arguments: arguments, + }) + } if err != nil { return "", fmt.Errorf("failed to call tool: %w", err) } @@ -166,6 +190,11 @@ func (a *Agent) callMCPTool(ctx context.Context, toolName string, arguments map[ func (a *Agent) ChatWithResponse(ctx context.Context, t *testing.T, userMessage string, expectedOntapErrorStr string) (string, error) { messages := []openai.ChatCompletionMessageParamUnion{ + openai.SystemMessage("When a tool call fails due to a validation error, analyze the error message and check if you " + + "misassigned a value to the wrong parameter (e.g. passed the cluster name as the svm name). " + + "Correct only the parameter mapping — never modify, truncate, or invent the actual values " + + "the user provided. If the error cannot be fixed by correcting a parameter mapping, report the " + + "error to the user and stop."), openai.UserMessage(userMessage), } From 5a785e56746f4970d4bbb6796457e668997bf653 Mon Sep 17 00:00:00 2001 From: rahulguptajss Date: Wed, 22 Apr 2026 13:07:05 +0530 Subject: [PATCH 2/5] ci: add session handling to agent --- integration/test/tools_test.go | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/integration/test/tools_test.go b/integration/test/tools_test.go index f125305..e140bb6 100644 --- a/integration/test/tools_test.go +++ b/integration/test/tools_test.go @@ -7,19 +7,18 @@ import ( "encoding/json" "errors" "fmt" - "log/slog" - "net/http" - "os" - "strings" - "testing" - "time" - "github.com/carlmjohnson/requests" "github.com/modelcontextprotocol/go-sdk/mcp" "github.com/netapp/ontap-mcp/config" "github.com/netapp/ontap-mcp/ontap" "github.com/openai/openai-go/v3" "github.com/openai/openai-go/v3/option" + "log/slog" + "net/http" + "os" + "strings" + "testing" + "time" ) const ( @@ -154,7 +153,7 @@ func (a *Agent) callMCPTool(ctx context.Context, toolName string, arguments map[ Name: toolName, Arguments: arguments, }) - if err != nil && strings.Contains(err.Error(), "connection closed") { + if err != nil && errors.Is(err, mcp.ErrConnectionClosed) { slog.Warn("MCP session dropped, reconnecting", slog.Any("error", err)) if reconnErr := a.reconnect(ctx); reconnErr != nil { return "", fmt.Errorf("failed to call tool: %w (reconnect failed: %w)", err, reconnErr) From c05b14eb321a1ad3b097abd0dea777ece4a20bd7 Mon Sep 17 00:00:00 2001 From: rahulguptajss Date: Wed, 22 Apr 2026 13:13:58 +0530 Subject: [PATCH 3/5] ci: add session handling to agent --- integration/test/tools_test.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/integration/test/tools_test.go b/integration/test/tools_test.go index e140bb6..980680b 100644 --- a/integration/test/tools_test.go +++ b/integration/test/tools_test.go @@ -138,7 +138,10 @@ func (a *Agent) convertMCPToolsToOpenAI() []openai.ChatCompletionToolUnionParam } func (a *Agent) reconnect(ctx context.Context) error { - _ = a.mcpSession.Close() + if a.mcpSession != nil { + _ = a.mcpSession.Close() + a.mcpSession = nil + } session, err := a.mcpClient.Connect(ctx, a.mcpTransport, nil) if err != nil { return fmt.Errorf("failed to reconnect to MCP server: %w", err) From 40f4c06c60ab78ae62f7ea7214ab6dd10dd6d2ba Mon Sep 17 00:00:00 2001 From: rahulguptajss Date: Wed, 22 Apr 2026 13:19:55 +0530 Subject: [PATCH 4/5] ci: add session handling to agent --- integration/test/tools_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/integration/test/tools_test.go b/integration/test/tools_test.go index 980680b..ce86564 100644 --- a/integration/test/tools_test.go +++ b/integration/test/tools_test.go @@ -159,7 +159,7 @@ func (a *Agent) callMCPTool(ctx context.Context, toolName string, arguments map[ if err != nil && errors.Is(err, mcp.ErrConnectionClosed) { slog.Warn("MCP session dropped, reconnecting", slog.Any("error", err)) if reconnErr := a.reconnect(ctx); reconnErr != nil { - return "", fmt.Errorf("failed to call tool: %w (reconnect failed: %w)", err, reconnErr) + return "", fmt.Errorf("failed to call tool: %w", errors.Join(err, fmt.Errorf("reconnect failed: %w", reconnErr))) } result, err = a.mcpSession.CallTool(ctx, &mcp.CallToolParams{ Name: toolName, From 4de6a0d2dd128f2b008034349fea38fde657f3dd Mon Sep 17 00:00:00 2001 From: rahulguptajss Date: Wed, 22 Apr 2026 15:56:46 +0530 Subject: [PATCH 5/5] ci: add session handling to agent --- integration/test/tools_test.go | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/integration/test/tools_test.go b/integration/test/tools_test.go index ce86564..4355547 100644 --- a/integration/test/tools_test.go +++ b/integration/test/tools_test.go @@ -192,11 +192,12 @@ func (a *Agent) callMCPTool(ctx context.Context, toolName string, arguments map[ func (a *Agent) ChatWithResponse(ctx context.Context, t *testing.T, userMessage string, expectedOntapErrorStr string) (string, error) { messages := []openai.ChatCompletionMessageParamUnion{ - openai.SystemMessage("When a tool call fails due to a validation error, analyze the error message and check if you " + - "misassigned a value to the wrong parameter (e.g. passed the cluster name as the svm name). " + - "Correct only the parameter mapping — never modify, truncate, or invent the actual values " + - "the user provided. If the error cannot be fixed by correcting a parameter mapping, report the " + - "error to the user and stop."), + openai.SystemMessage("When a tool call fails due to a validation error, check two things: " + + "(1) Did you misassign a value to the wrong parameter (e.g. passed the cluster name as the svm name)? " + + "(2) Did you misread or partially extract a value from the user's message (e.g. dropped a path prefix like /vol/vol1/)? " + + "In either case, correct the mistake using only values already present in the conversation — " + + "never truncate, invent, or change values beyond what is needed to fix the extraction or mapping error. " + + "If the error cannot be fixed this way, report it to the user and stop."), openai.UserMessage(userMessage), }