chore: [ML-1395]: Updated download logs tool to accept prefix argument (#191)

* fb760e lets call it log key wherever possible (#210)

* d41a28 chore: [ML-1395]: Resolved comments

* 210394 chore: [ML-1395]: Resolved comments

* 71fbee chore: [ML-1395]: Resolved comments

* e40cd3 chore: [ML-1395]: Resolved conflicts

* 1de7af chore: [ML-1395]: update gci version

* ff71a0 Merge branch 'master' of https://git0.harness.io/l7B_kbSEQD2wjrM7PShm5w/PROD/Harness_Commons/mcp-server into ML-1395

* 85817c chore: [ML-1395]: Updated download logs tool to accept prefix argument

Author: Saranya-jena

client/logs.go

@@ -3,6 +3,7 @@ package client
 import (
 	"context"
 	"fmt"
+	"log/slog"
 	"time"
 
 	"github.com/harness/harness-mcp/client/dto"
@@ -19,37 +20,46 @@ type LogService struct {
 }
 
 // DownloadLogs fetches a download URL for pipeline execution logs
-func (l *LogService) DownloadLogs(ctx context.Context, scope dto.Scope, planExecutionID string) (string, error) {
-	// First, get the pipeline execution details to determine the prefix format
-	pipelineService := &PipelineService{Client: l.PipelineClient} // TODO: needs to be changed for internal case, we should move this above
-	execution, err := pipelineService.GetExecution(ctx, scope, planExecutionID)
-	if err != nil {
-		return "", fmt.Errorf("failed to get execution details: %w", err)
-	}
-
-	// Build the prefix based on the execution details
-	var prefix string
-	if execution.Data.ShouldUseSimplifiedBaseKey {
-		// Simplified key format
-		prefix = fmt.Sprintf("%s/pipeline/%s/%d/-%s",
-			scope.AccountID,
-			execution.Data.PipelineIdentifier,
-			execution.Data.RunSequence,
-			planExecutionID)
+// If logKey is not empty, it will use that log key to fetch logs instead of building one from execution details
+func (l *LogService) DownloadLogs(ctx context.Context, scope dto.Scope, planExecutionID string, logKey string) (string, error) {
+	// Use custom log key if provided, otherwise build it from execution details
+	var finalLogKey string
+	var err error
+	if logKey != "" {
+		slog.Info("Using custom log key for log download", "logKey", logKey)
+		finalLogKey = logKey
 	} else {
-		// Standard key format
-		prefix = fmt.Sprintf("accountId:%s/orgId:%s/projectId:%s/pipelineId:%s/runSequence:%d/level0:pipeline",
-			scope.AccountID,
-			execution.Data.OrgIdentifier,
-			execution.Data.ProjectIdentifier,
-			execution.Data.PipelineIdentifier,
-			execution.Data.RunSequence)
+		slog.Info("Building log key for log download from execution details")
+		// First, get the pipeline execution details to determine the prefix format
+		pipelineService := &PipelineService{Client: l.PipelineClient} // TODO: needs to be changed for internal case, we should move this above
+		execution, err := pipelineService.GetExecution(ctx, scope, planExecutionID)
+		if err != nil {
+			return "", fmt.Errorf("failed to get execution details: %w", err)
+		}
+
+		// Build the log key based on the execution details
+		if execution.Data.ShouldUseSimplifiedBaseKey {
+			// Simplified key format
+			finalLogKey = fmt.Sprintf("%s/pipeline/%s/%d/-%s",
+				scope.AccountID,
+				execution.Data.PipelineIdentifier,
+				execution.Data.RunSequence,
+				planExecutionID)
+		} else {
+			// Standard key format
+			finalLogKey = fmt.Sprintf("accountId:%s/orgId:%s/projectId:%s/pipelineId:%s/runSequence:%d/level0:pipeline",
+				scope.AccountID,
+				execution.Data.OrgIdentifier,
+				execution.Data.ProjectIdentifier,
+				execution.Data.PipelineIdentifier,
+				execution.Data.RunSequence)
+		}
 	}
 
 	// Prepare query parameters
 	params := make(map[string]string)
 	params["accountID"] = scope.AccountID
-	params["prefix"] = prefix
+	params["prefix"] = finalLogKey
 
 	// Initialize the response object
 	response := &dto.LogDownloadResponse{}

pkg/harness/tools/logs.go

@@ -22,13 +22,15 @@ func DownloadExecutionLogsTool(config *config.Config, client *client.LogService)
 	return mcp.NewTool("download_execution_logs",
 			mcp.WithDescription("Downloads logs for an execution inside Harness"),
 			mcp.WithString("plan_execution_id",
-				mcp.Required(),
 				mcp.Description("The ID of the plan execution"),
 			),
 			mcp.WithString("logs_directory",
 				mcp.Required(),
 				mcp.Description("The absolute path to the directory where the logs should get downloaded"),
 			),
+			mcp.WithString("log_key",
+				mcp.Description("Optional log key to be used for downloading logs directly"),
+			),
 			common.WithScope(config, true),
 		),
 		func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
@@ -42,7 +44,10 @@ func DownloadExecutionLogsTool(config *config.Config, client *client.LogService)
 				return mcp.NewToolResultError(err.Error()), nil
 			}
 
-			logDownloadURL, err := client.DownloadLogs(ctx, scope, planExecutionID)
+			// Get optional log key parameter
+			logKey, _ := OptionalParam[string](request, "log_key")
+
+			logDownloadURL, err := client.DownloadLogs(ctx, scope, planExecutionID, logKey)
 			if err != nil {
 				return nil, fmt.Errorf("failed to fetch log download URL: %w", err)
 			}
@@ -76,7 +81,7 @@ func DownloadExecutionLogsTool(config *config.Config, client *client.LogService)
 			}
 
 			// Get the download URL
-			logDownloadURL, err = client.DownloadLogs(ctx, scope, planExecutionID)
+			logDownloadURL, err = client.DownloadLogs(ctx, scope, planExecutionID, logKey)
 			if err != nil {
 				return mcp.NewToolResultError(fmt.Sprintf("failed to fetch log download URL: %v", err)), nil
 			}