1
if !s.EnableToolUseShim {
2
    var functionDefinitions []*gollm.FunctionDefinition
3
    for _, tool := range s.Tools.AllTools() {
4
        functionDefinitions = append(functionDefinitions, tool.FunctionDefinition())
5
    }
6
    // Sort function definitions to help KV cache reuse
7
    sort.Slice(functionDefinitions, func(i, j int) bool {
8
        return functionDefinitions[i].Name < functionDefinitions[j].Name
9
    })
10
    if err := s.llmChat.SetFunctionDefinitions(functionDefinitions); err != nil {
11
        return fmt.Errorf("setting function definitions: %w", err)
12
    }
13
}

1
package main
2

3
import (
4
  "context"
5
  "fmt"
6
  "github.com/GoogleCloudPlatform/kubectl-ai/pkg/tools"
7
  "github.com/mark3labs/mcp-go/mcp"
8
  "github.com/mark3labs/mcp-go/server"
9
  "k8s.io/klog/v2"
10
)
11

12
type kubectlMCPServer struct {
13
  kubectlConfig string
14
  server        *server.MCPServer
15
  tools         tools.Tools
16
  workDir       string
17
}
18

19
func newKubectlMCPServer(ctx context.Context, kubectlConfig string, tools tools.Tools, workDir string) (*kubectlMCPServer, error) {
20
  s := &kubectlMCPServer{
21
    kubectlConfig: kubectlConfig,
22
    workDir:       workDir,
23
    server: server.NewMCPServer(
24
      "kubectl-ai",
25
      "0.0.1",
26
      server.WithToolCapabilities(true),
27
    ),
28
    tools: tools,
29
  }
30
  for _, tool := range s.tools.AllTools() {
31
    toolDefn := tool.FunctionDefinition()
32
    toolInputSchema, err := toolDefn.Parameters.ToRawSchema()
33
    if err != nil {
34
      return nil, fmt.Errorf("converting tool schema to json.RawMessage: %w", err)
35
    }
36
    s.server.AddTool(mcp.NewToolWithRawSchema(
37
      toolDefn.Name,
38
      toolDefn.Description,
39
      toolInputSchema,
40
    ), s.handleToolCall)
41
  }
42
  return s, nil
43
}
44
func (s *kubectlMCPServer) Serve(ctx context.Context) error {
45
  return server.ServeStdio(s.server)
46
}
47

48
func (s *kubectlMCPServer) handleToolCall(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
49

50
  log := klog.FromContext(ctx)
51

52
  name := request.Params.Name
53
  command := request.Params.Arguments["command"].(string)
54
  modifiesResource := request.Params.Arguments["modifies_resource"].(string)
55
  log.Info("Received tool call", "tool", name, "command", command, "modifies_resource", modifiesResource)
56

57
  ctx = context.WithValue(ctx, tools.KubeconfigKey, s.kubectlConfig)
58
  ctx = context.WithValue(ctx, tools.WorkDirKey, s.workDir)
59

60
  tool := tools.Lookup(name)
61
  if tool == nil {
62
    return &mcp.CallToolResult{
63
      Content: []mcp.Content{
64
        mcp.TextContent{
65
          Type: "text",
66
          Text: fmt.Sprintf("Error: Tool %s not found", name),
67
        },
68
      },
69
    }, nil
70
  }
71
  output, err := tool.Run(ctx, map[string]any{
72
    "command": command,
73
  })
74
  if err != nil {
75
    log.Error(err, "Error running tool call")
76
    return &mcp.CallToolResult{
77
      Content: []mcp.Content{
78
        mcp.TextContent{
79
          Type: "text",
80
          Text: fmt.Sprintf("Error: %v", err),
81
        },
82
      },
83
      IsError: true,
84
    }, nil
85
  }
86

87
  result, err := tools.ToolResultToMap(output)
88
  if err != nil {
89
    log.Error(err, "Error converting tool call output to result")
90
    return &mcp.CallToolResult{
91
      Content: []mcp.Content{
92
        mcp.TextContent{
93
          Type: "text",
94
          Text: fmt.Sprintf("Error: %v", err),
95
        },
96
      },
97
      IsError: true,
98
    }, nil
99
  }
100

101
  log.Info("Tool call output", "tool", name, "result", result)
102

103
  return &mcp.CallToolResult{
104
    Content: []mcp.Content{
105
      mcp.TextContent{
106
        Type: "text",
107
        Text: fmt.Sprintf("%v", result),
108
      },
109
    },
110
  }, nil
111
}