feat: [ML-1197]: Improve reliability of the uber agent (#41)

* feat: [ML-1197]: Fix IDP headers

* feat: [ML-1197]: Add unit tests for client

* feat: [ML-1197]: make format

* feat: [ML-1197]: Fix dashboard unmarshalling error and 422 from ai-devops

Author: David-Kreiner

client/client.go

@@ -355,6 +355,13 @@ func unmarshalResponse(resp *http.Response, data interface{}) error {
 		// If we couldn't parse it as an error response, continue with normal unmarshaling
 	}
 
+	// Special handling for string responses: if data is a string pointer, directly assign the body content
+	if strPtr, ok := data.(*string); ok {
+		*strPtr = string(body)
+		return nil
+	}
+
+	// Otherwise try to unmarshal as JSON
 	err = json.Unmarshal(body, data)
 	if err != nil {
 		return fmt.Errorf("error deserializing response body : %w - original response: %s", err, string(body))

client/client_test.go

@@ -0,0 +1,241 @@
+package client
+
+import (
+	"bytes"
+	"errors"
+	"io"
+	"net/http"
+	"strings"
+	"testing"
+)
+
+// TestUnmarshalResponse_StringPointer tests the string pointer handling in unmarshalResponse function
+func TestUnmarshalResponse_StringPointer(t *testing.T) {
+	tests := []struct {
+		name           string
+		responseBody   string
+		expectedResult string
+	}{
+		{
+			name:           "Plain text response",
+			responseBody:   "This is a plain text response",
+			expectedResult: "This is a plain text response",
+		},
+		{
+			name:           "JSON-like text response",
+			responseBody:   "{\"key\": \"value\"}",
+			expectedResult: "{\"key\": \"value\"}",
+		},
+		{
+			name:           "Empty response",
+			responseBody:   "",
+			expectedResult: "",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create a mock response
+			resp := &http.Response{
+				Body: io.NopCloser(bytes.NewBufferString(tt.responseBody)),
+			}
+
+			// Create a string pointer to unmarshal into
+			var result string
+
+			// Call the function
+			err := unmarshalResponse(resp, &result)
+
+			// Assert no error occurred
+			if err != nil {
+				t.Errorf("Expected no error, got %v", err)
+			}
+
+			// Assert the result matches expected
+			if tt.expectedResult != result {
+				t.Errorf("Expected result %v, got %v", tt.expectedResult, result)
+			}
+		})
+	}
+}
+
+// TestUnmarshalResponse_JSONStruct tests the JSON unmarshaling into struct types
+func TestUnmarshalResponse_JSONStruct(t *testing.T) {
+	// Define a test struct that matches a sample JSON response
+	type testStruct struct {
+		Name    string `json:"name"`
+		Value   int    `json:"value"`
+		Enabled bool   `json:"enabled"`
+	}
+
+	tests := []struct {
+		name           string
+		responseBody   string
+		expectedResult testStruct
+	}{
+		{
+			name:         "Valid JSON struct",
+			responseBody: `{"name":"test","value":42,"enabled":true}`,
+			expectedResult: testStruct{
+				Name:    "test",
+				Value:   42,
+				Enabled: true,
+			},
+		},
+		{
+			name:         "Partial JSON struct",
+			responseBody: `{"name":"partial"}`,
+			expectedResult: testStruct{
+				Name:    "partial",
+				Value:   0,
+				Enabled: false,
+			},
+		},
+		{
+			name:         "Complex nested JSON",
+			responseBody: `{"name":"nested","value":123,"enabled":true}`,
+			expectedResult: testStruct{
+				Name:    "nested",
+				Value:   123,
+				Enabled: true,
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create a mock response
+			resp := &http.Response{
+				Body: io.NopCloser(bytes.NewBufferString(tt.responseBody)),
+			}
+
+			// Create a struct to unmarshal into
+			var result testStruct
+
+			// Call the function
+			err := unmarshalResponse(resp, &result)
+
+			// Assert no error occurred
+			if err != nil {
+				t.Errorf("Expected no error, got %v", err)
+			}
+
+			// Assert the result matches expected
+			if tt.expectedResult.Name != result.Name ||
+				tt.expectedResult.Value != result.Value ||
+				tt.expectedResult.Enabled != result.Enabled {
+				t.Errorf("Expected result %+v, got %+v", tt.expectedResult, result)
+			}
+		})
+	}
+}
+
+// TestUnmarshalResponse_ErrorResponse tests the error response handling in unmarshalResponse function
+func TestUnmarshalResponse_ErrorResponse(t *testing.T) {
+	tests := []struct {
+		name          string
+		responseBody  string
+		statusCode    int
+		expectedError string
+	}{
+		{
+			name:          "API error response",
+			responseBody:  `{"code":"INVALID_REQUEST","message":"Invalid input parameters"}`,
+			statusCode:    400,
+			expectedError: "API error: Invalid input parameters",
+		},
+		{
+			name:         "Non-JSON error response",
+			responseBody: "Internal server error occurred",
+			statusCode:   500,
+			// Non-JSON content will cause unmarshal error
+			expectedError: "error deserializing response body",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create a mock response with error status code
+			resp := &http.Response{
+				StatusCode: tt.statusCode,
+				Body:       io.NopCloser(bytes.NewBufferString(tt.responseBody)),
+			}
+
+			// Create a struct to unmarshal into
+			var result map[string]interface{}
+
+			// Call the function
+			err := unmarshalResponse(resp, &result)
+
+			// Assert error occurred and contains expected message
+			if err == nil {
+				t.Errorf("Expected error, got nil")
+			} else if !strings.Contains(err.Error(), tt.expectedError) {
+				t.Errorf("Expected error containing '%s', got '%s'", tt.expectedError, err.Error())
+			}
+		})
+	}
+}
+
+// TestUnmarshalResponse_ErrorCases tests various error handling scenarios in unmarshalResponse function
+func TestUnmarshalResponse_ErrorCases(t *testing.T) {
+	tests := []struct {
+		name          string
+		resp          *http.Response
+		expectedError string
+	}{
+		{
+			name:          "Nil response",
+			resp:          nil,
+			expectedError: "http response body is not available",
+		},
+		{
+			name: "Nil body",
+			resp: &http.Response{
+				Body: nil,
+			},
+			expectedError: "http response body is not available",
+		},
+		{
+			name: "Read body error",
+			resp: &http.Response{
+				Body: &errorReader{},
+			},
+			expectedError: "error reading response body",
+		},
+		{
+			name: "Invalid JSON",
+			resp: &http.Response{
+				Body: io.NopCloser(bytes.NewBufferString("{invalid json:}")),
+			},
+			expectedError: "error deserializing response body",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create a struct to unmarshal into
+			var result map[string]interface{}
+
+			// Call the function
+			err := unmarshalResponse(tt.resp, &result)
+
+			// Assert error occurred and contains expected message
+			if err == nil {
+				t.Errorf("Expected error, got nil")
+			} else if !strings.Contains(err.Error(), tt.expectedError) {
+				t.Errorf("Expected error containing '%s', got '%s'", tt.expectedError, err.Error())
+			}
+		})
+	}
+}
+
+type errorReader struct{}
+
+func (e *errorReader) Read(p []byte) (n int, err error) {
+	return 0, errors.New("mock read error")
+}
+
+func (e *errorReader) Close() error {
+	return nil
+}

client/dto/genai.go

@@ -1,13 +1,19 @@
 package dto
 
+type ContextType string
+
+const (
+	ContextTypeOther ContextType = "other"
+)
+
 type Capability struct {
 	Type    string `json:"type"`
 	Version string `json:"version"`
 }
 
 type ContextItem struct {
-	Type    string `json:"type"`
-	Payload any    `json:"payload"`
+	Type    ContextType `json:"type"`
+	Payload any         `json:"payload"`
 }
 
 type HarnessContext struct {

client/idp.go

@@ -45,7 +45,7 @@ func (i *IDPService) GetEntity(ctx context.Context, scope dto.Scope, kind string
 
 	response := new(dto.EntityResponse)
 
-	err := i.Client.Get(ctx, path, params, map[string]string{}, response)
+	err := i.Client.Get(ctx, path, params, headers, response)
 	if err != nil {
 		return nil, err
 	}
@@ -109,7 +109,7 @@ func (i *IDPService) ListEntities(ctx context.Context, scope dto.Scope, getEntit
 
 	response := make([]dto.EntityResponse, 0)
 
-	err := i.Client.Get(ctx, path, params, map[string]string{}, &response)
+	err := i.Client.Get(ctx, path, params, headers, &response)
 	if err != nil {
 		return nil, err
 	}

pkg/harness/genai.go

@@ -38,10 +38,12 @@ func AIDevOpsAgentTool(config *config.Config, client *client.GenaiService) (tool
 					"properties": map[string]any{
 						"type": map[string]any{
 							"type":        "string",
-							"description": "The type of context item",
+							"description": "The type of context item (other)",
+							"enum":        []string{string(dto.ContextTypeOther)},
 						},
 						"payload": map[string]any{
-							"description": "The payload for this context item",
+							"type":        []string{"object", "array", "string", "number", "boolean"},
+							"description": "The payload for this context item, accepts any valid JSON value. Example: {\"stage_type\": \"Custom\"}",
 						},
 					},
 					"required": []string{"type", "payload"},
@@ -112,7 +114,7 @@ func AIDevOpsAgentTool(config *config.Config, client *client.GenaiService) (tool
 					ctxType, _ := ctxMap["type"].(string)
 					ctxPayload := ctxMap["payload"]
 					contextItems = append(contextItems, dto.ContextItem{
-						Type:    ctxType,
+						Type:    dto.ContextType(ctxType),
 						Payload: ctxPayload,
 					})
 				}

pkg/harness/pipelines.go

@@ -14,7 +14,7 @@ import (
 
 func GetPipelineTool(config *config.Config, client *client.PipelineService) (tool mcp.Tool, handler server.ToolHandlerFunc) {
 	return mcp.NewTool("get_pipeline",
-			mcp.WithDescription("Get details of a specific pipeline in a Harness repository."),
+			mcp.WithDescription("Get details of a specific pipeline in a Harness repository. Use list_pipelines (if available) first to find the correct pipeline_id if you're unsure of the exact ID."),
 			mcp.WithString("pipeline_id",
 				mcp.Required(),
 				mcp.Description("The ID of the pipeline"),

pkg/harness/tools.go

@@ -25,7 +25,7 @@ const aiServiceIdentity = "aifoundation"
 var defaultJWTLifetime = 1 * time.Hour
 
 // Default timeout for GenAI service
-const defaultGenaiTimeout = 60 * time.Second
+const defaultGenaiTimeout = 300 * time.Second
 
 // InitToolsets initializes and returns the toolset groups
 func InitToolsets(config *config.Config) (*toolsets.ToolsetGroup, error) {