From 7dff2fbb4e0ab0f5575aa74d5dbcbb701461fa8b Mon Sep 17 00:00:00 2001
From: Karel Suta <ksuta@redhat.com>
Date: Thu, 10 Apr 2025 23:48:24 +0200
Subject: [PATCH] feat: add optional container parameter for pods_log tool

---
 pkg/kubernetes/pods.go |  3 ++-
 pkg/mcp/pods.go        |  7 ++++++-
 pkg/mcp/pods_test.go   | 30 ++++++++++++++++++++++++++++++
 3 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/pkg/kubernetes/pods.go b/pkg/kubernetes/pods.go
index ada91826..6a728a8b 100644
--- a/pkg/kubernetes/pods.go
+++ b/pkg/kubernetes/pods.go
@@ -77,10 +77,11 @@ func (k *Kubernetes) PodsDelete(ctx context.Context, namespace, name string) (st
 		k.clientSet.CoreV1().Pods(namespace).Delete(ctx, name, metav1.DeleteOptions{})
 }
 
-func (k *Kubernetes) PodsLog(ctx context.Context, namespace, name string) (string, error) {
+func (k *Kubernetes) PodsLog(ctx context.Context, namespace, name, container string) (string, error) {
 	tailLines := int64(256)
 	req := k.clientSet.CoreV1().Pods(namespaceOrDefault(namespace)).GetLogs(name, &v1.PodLogOptions{
 		TailLines: &tailLines,
+		Container: container,
 	})
 	res := req.Do(ctx)
 	if res.Error() != nil {
diff --git a/pkg/mcp/pods.go b/pkg/mcp/pods.go
index 8508d690..0d3f72a0 100644
--- a/pkg/mcp/pods.go
+++ b/pkg/mcp/pods.go
@@ -49,6 +49,7 @@ func (s *Server) initPods() []server.ServerTool {
 			mcp.WithDescription("Get the logs of a Kubernetes Pod in the current or provided namespace with the provided name"),
 			mcp.WithString("namespace", mcp.Description("Namespace to get the Pod logs from")),
 			mcp.WithString("name", mcp.Description("Name of the Pod to get the logs from"), mcp.Required()),
+			mcp.WithString("container", mcp.Description("Name of the Pod container to get the logs from (Optional)")),
 		), s.podsLog},
 		{mcp.NewTool("pods_run",
 			mcp.WithDescription("Run a Kubernetes Pod in the current or provided namespace with the provided container image and optional name"),
@@ -150,7 +151,11 @@ func (s *Server) podsLog(ctx context.Context, ctr mcp.CallToolRequest) (*mcp.Cal
 	if name == nil {
 		return NewTextResult("", errors.New("failed to get pod log, missing argument name")), nil
 	}
-	ret, err := s.k.PodsLog(ctx, ns.(string), name.(string))
+	container := ctr.Params.Arguments["container"]
+	if container == nil {
+		container = ""
+	}
+	ret, err := s.k.PodsLog(ctx, ns.(string), name.(string), container.(string))
 	if err != nil {
 		return NewTextResult("", fmt.Errorf("failed to get pod %s log in namespace %s: %v", name, ns, err)), nil
 	} else if ret == "" {
diff --git a/pkg/mcp/pods_test.go b/pkg/mcp/pods_test.go
index b7ca84e1..5c2898ec 100644
--- a/pkg/mcp/pods_test.go
+++ b/pkg/mcp/pods_test.go
@@ -526,6 +526,36 @@ func TestPodsLog(t *testing.T) {
 				return
 			}
 		})
+		podsContainerLogInNamespace, err := c.callTool("pods_log", map[string]interface{}{
+			"namespace": "ns-1",
+			"name":      "a-pod-in-ns-1",
+			"container": "nginx",
+		})
+		t.Run("pods_log with name, container and namespace returns pod log", func(t *testing.T) {
+			if err != nil {
+				t.Fatalf("call tool failed %v", err)
+				return
+			}
+			if podsContainerLogInNamespace.IsError {
+				t.Fatalf("call tool failed")
+				return
+			}
+		})
+		toolResult, err := c.callTool("pods_log", map[string]interface{}{
+			"namespace": "ns-1",
+			"name":      "a-pod-in-ns-1",
+			"container": "a-not-existing-container",
+		})
+		t.Run("pods_log with non existing container returns error", func(t *testing.T) {
+			if toolResult.IsError != true {
+				t.Fatalf("call tool should fail")
+				return
+			}
+			if toolResult.Content[0].(mcp.TextContent).Text != "failed to get pod a-pod-in-ns-1 log in namespace ns-1: container a-not-existing-container is not valid for pod a-pod-in-ns-1" {
+				t.Fatalf("invalid error message, got %v", toolResult.Content[0].(mcp.TextContent).Text)
+				return
+			}
+		})
 	})
 }