Skip to content

Commit d0deb3f

Browse files
Saranya-jenaHarness
authored andcommitted
fix: [ML-1335]: Pass harness-account header in API calls (#150)
* Merge remote-tracking branch 'origin/master' into ML-1335 * revert change to templates.go take up in new PR * re-create pullreq.go * re-create accesscontrol.go * make format * fix compile * Merge remote-tracking branch 'origin/master' into ML-1335 * code format + remove check on account_id * merge scope_utils.go into scope.go * remove context.go * port left over calls to scope methods * Merge remote-tracking branch 'origin' into ML-1335 * Set account ID in context using middleware (#160) * fix: [ML-1335]: fixed e2e Signed-off-by: Saranya-jena <[email protected]> * add unit tests * remove AccountIDKey * rename to "scope" * Set account ID in context using middleware * fix: [ML-1335]: Resolved comments Signed-off-by: Saranya-jena <[email protected]> * fix: [ML-1335]: fixed 404 issue with list_templates tool Signed-off-by: Saranya-jena <[email protected]>
1 parent d57b791 commit d0deb3f

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

41 files changed

+577
-211
lines changed

client/client.go

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"github.com/cenkalti/backoff/v4"
1616
"github.com/harness/harness-mcp/client/dto"
1717
"github.com/harness/harness-mcp/pkg/harness/auth"
18+
"github.com/harness/harness-mcp/pkg/harness/common"
1819
)
1920

2021
const (
@@ -29,8 +30,6 @@ var (
2930
// when different tools get added.
3031
defaultPageSize = 5
3132
maxPageSize = 20
32-
33-
apiKeyHeader = "x-api-key"
3433
)
3534

3635
var (
@@ -475,6 +474,13 @@ func (c *Client) Do(r *http.Request) (*http.Response, error) {
475474
}
476475
r.Header.Set(k, v)
477476

477+
// Check for scope in context and add account ID to headers if present
478+
if scope, err := common.GetScopeFromContext(ctx); err == nil {
479+
if scope.AccountID != "" {
480+
slog.Debug("Adding account ID header from scope in context", "accountID", scope.AccountID)
481+
r.Header.Set("Harness-Account", scope.AccountID)
482+
}
483+
}
478484
return c.client.Do(r)
479485
}
480486

client/client_test.go

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,15 @@ package client
22

33
import (
44
"bytes"
5+
"context"
56
"errors"
67
"io"
78
"net/http"
89
"strings"
910
"testing"
11+
12+
"github.com/harness/harness-mcp/client/dto"
13+
"github.com/harness/harness-mcp/pkg/harness/common"
1014
)
1115

1216
// TestUnmarshalResponse_StringPointer tests the string pointer handling in unmarshalResponse function
@@ -239,3 +243,75 @@ func (e *errorReader) Read(p []byte) (n int, err error) {
239243
func (e *errorReader) Close() error {
240244
return nil
241245
}
246+
247+
// --- New tests for Harness-Account header injection via scope ---
248+
249+
type mockAuthProvider struct{}
250+
251+
func (m mockAuthProvider) GetHeader(ctx context.Context) (string, string, error) {
252+
return "Authorization", "Bearer test-token", nil
253+
}
254+
255+
type recordingTransport struct {
256+
sawAccount string
257+
}
258+
259+
func (rt *recordingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
260+
rt.sawAccount = req.Header.Get("Harness-Account")
261+
return &http.Response{
262+
StatusCode: http.StatusOK,
263+
Body: io.NopCloser(strings.NewReader("ok")),
264+
Header: make(http.Header),
265+
Request: req,
266+
}, nil
267+
}
268+
269+
func TestClientDo_AddsHarnessAccountHeader_FromScope(t *testing.T) {
270+
// Prepare client with recording transport
271+
rt := &recordingTransport{}
272+
c := &Client{
273+
client: &http.Client{Transport: rt},
274+
AuthProvider: mockAuthProvider{},
275+
}
276+
277+
// Build request with scope in context
278+
req, err := http.NewRequest(http.MethodGet, "http://example.com/test", nil)
279+
if err != nil {
280+
t.Fatalf("failed to create request: %v", err)
281+
}
282+
283+
scope := dto.Scope{AccountID: "acc_123", OrgID: "org", ProjectID: "proj"}
284+
ctx := common.WithScopeContext(context.Background(), scope)
285+
req = req.WithContext(ctx)
286+
287+
// Execute
288+
if _, err := c.Do(req); err != nil {
289+
t.Fatalf("Do() returned error: %v", err)
290+
}
291+
292+
// Assert
293+
if rt.sawAccount != scope.AccountID {
294+
t.Fatalf("expected Harness-Account header %q, got %q", scope.AccountID, rt.sawAccount)
295+
}
296+
}
297+
298+
func TestClientDo_DoesNotAddHarnessAccountHeader_WhenNoScope(t *testing.T) {
299+
rt := &recordingTransport{}
300+
c := &Client{
301+
client: &http.Client{Transport: rt},
302+
AuthProvider: mockAuthProvider{},
303+
}
304+
305+
req, err := http.NewRequest(http.MethodGet, "http://example.com/test", nil)
306+
if err != nil {
307+
t.Fatalf("failed to create request: %v", err)
308+
}
309+
310+
if _, err := c.Do(req); err != nil {
311+
t.Fatalf("Do() returned error: %v", err)
312+
}
313+
314+
if rt.sawAccount != "" {
315+
t.Fatalf("expected no Harness-Account header, got %q", rt.sawAccount)
316+
}
317+
}

client/dto/intelligence.go

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -107,14 +107,14 @@ func (s *ServiceChatParameters) IsStreaming() bool {
107107
}
108108

109109
type ServiceChatResponseIntelligence struct {
110-
ConversationID string `json:"conversation_id"`
111-
ConversationRaw string `json:"conversation_raw"`
112-
InteractionID string `json:"interaction_id,omitempty"`
113-
CapabilitiesToRun []map[string]interface{} `json:"capabilities_to_run,omitempty"`
114-
ModelUsage map[string]interface{} `json:"model_usage,omitempty"`
115-
// Keep these for backward compatibility if needed elsewhere
116-
Response string `json:"response,omitempty"`
117-
Error string `json:"error,omitempty"`
110+
ConversationID string `json:"conversation_id"`
111+
ConversationRaw string `json:"conversation_raw"`
112+
InteractionID string `json:"interaction_id,omitempty"`
113+
CapabilitiesToRun []map[string]interface{} `json:"capabilities_to_run,omitempty"`
114+
ModelUsage map[string]interface{} `json:"model_usage,omitempty"`
115+
// Keep these for backward compatibility if needed elsewhere
116+
Response string `json:"response,omitempty"`
117+
Error string `json:"error,omitempty"`
118118
}
119119

120120
// ProgressUpdate represents progress information for streaming responses

cmd/harness-mcp-server/main.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -366,7 +366,7 @@ func runStdioServer(ctx context.Context, config config.Config) error {
366366

367367
// Create server
368368
// WithRecovery makes sure panics are logged and don't crash the server
369-
harnessServer := harness.NewServer(version, server.WithHooks(hooks), server.WithRecovery())
369+
harnessServer := harness.NewServer(version, &config, server.WithHooks(hooks), server.WithRecovery())
370370

371371
// Initialize toolsets
372372
toolsets, err := harness.InitToolsets(ctx, &config)

pkg/harness/tools/scope.go renamed to pkg/harness/common/scope.go

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
1-
package tools
1+
package common
22

33
import (
4+
"context"
45
"fmt"
56

67
"github.com/harness/harness-mcp/client/dto"
@@ -83,3 +84,37 @@ func FetchScope(config *config.Config, request mcp.CallToolRequest, required boo
8384

8485
return scope, nil
8586
}
87+
88+
// OptionalParam is a helper function that can be used to fetch an optional parameter from the request.
89+
func OptionalParam[T any](r mcp.CallToolRequest, p string) (T, error) {
90+
var zero T
91+
92+
// Check if the parameter is present in the request
93+
if _, ok := r.GetArguments()[p]; !ok {
94+
return zero, nil
95+
}
96+
97+
// Check if the parameter is of the expected type
98+
if _, ok := r.GetArguments()[p].(T); !ok {
99+
return zero, fmt.Errorf("parameter %s is not of type %T, is %T", p, zero, r.GetArguments()[p])
100+
}
101+
102+
return r.GetArguments()[p].(T), nil
103+
}
104+
105+
// ScopeKey is the context key for storing the scope
106+
type ScopeKey struct{}
107+
108+
// GetScopeFromContext retrieves the scope from the context
109+
func GetScopeFromContext(ctx context.Context) (dto.Scope, error) {
110+
scope, ok := ctx.Value(ScopeKey{}).(dto.Scope)
111+
if !ok {
112+
return dto.Scope{}, fmt.Errorf("scope not found in context")
113+
}
114+
return scope, nil
115+
}
116+
117+
// WithScopeContext adds the scope to the context
118+
func WithScopeContext(ctx context.Context, scope dto.Scope) context.Context {
119+
return context.WithValue(ctx, ScopeKey{}, scope)
120+
}

pkg/harness/common/scope_test.go

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,96 @@
1+
package common
2+
3+
import (
4+
"testing"
5+
6+
"github.com/harness/harness-mcp/cmd/harness-mcp-server/config"
7+
"github.com/mark3labs/mcp-go/mcp"
8+
)
9+
10+
func newRequest(args map[string]any) mcp.CallToolRequest {
11+
r := mcp.CallToolRequest{}
12+
r.Params.Arguments = args
13+
return r
14+
}
15+
16+
func TestFetchScope_ErrorWhenNoAccountID(t *testing.T) {
17+
cfg := &config.Config{AccountID: ""}
18+
r := newRequest(nil)
19+
_, err := FetchScope(cfg, r, false)
20+
if err == nil || err.Error() != "account ID is required" {
21+
t.Fatalf("expected account ID required error, got: %v", err)
22+
}
23+
}
24+
25+
func TestFetchScope_UsesConfigDefaultsWhenNotProvided(t *testing.T) {
26+
cfg := &config.Config{AccountID: "acc", DefaultOrgID: "org", DefaultProjectID: "proj"}
27+
r := newRequest(nil)
28+
scope, err := FetchScope(cfg, r, false)
29+
if err != nil {
30+
t.Fatalf("unexpected err: %v", err)
31+
}
32+
if scope.AccountID != "acc" || scope.OrgID != "org" || scope.ProjectID != "proj" {
33+
t.Fatalf("unexpected scope: %+v", scope)
34+
}
35+
}
36+
37+
func TestFetchScope_RequestOverridesConfig(t *testing.T) {
38+
cfg := &config.Config{AccountID: "acc", DefaultOrgID: "org", DefaultProjectID: "proj"}
39+
args := map[string]any{"org_id": "o2", "project_id": "p2"}
40+
r := newRequest(args)
41+
scope, err := FetchScope(cfg, r, false)
42+
if err != nil {
43+
t.Fatalf("unexpected err: %v", err)
44+
}
45+
if scope.OrgID != "o2" || scope.ProjectID != "p2" {
46+
t.Fatalf("expected overrides (o2,p2), got: (%s,%s)", scope.OrgID, scope.ProjectID)
47+
}
48+
}
49+
50+
func TestFetchScope_RequiredMissingOrg(t *testing.T) {
51+
cfg := &config.Config{AccountID: "acc", DefaultOrgID: "", DefaultProjectID: "proj"}
52+
r := newRequest(nil)
53+
_, err := FetchScope(cfg, r, true)
54+
if err == nil || err.Error() != "org ID is required" {
55+
t.Fatalf("expected org ID required error, got: %v", err)
56+
}
57+
}
58+
59+
func TestFetchScope_RequiredMissingProject(t *testing.T) {
60+
cfg := &config.Config{AccountID: "acc", DefaultOrgID: "org", DefaultProjectID: ""}
61+
r := newRequest(nil)
62+
_, err := FetchScope(cfg, r, true)
63+
if err == nil || err.Error() != "project ID is required" {
64+
t.Fatalf("expected project ID required error, got: %v", err)
65+
}
66+
}
67+
68+
func TestOptionalParam_MissingReturnsZero(t *testing.T) {
69+
r := newRequest(nil)
70+
val, err := OptionalParam[string](r, "missing")
71+
if err != nil {
72+
t.Fatalf("unexpected err: %v", err)
73+
}
74+
if val != "" {
75+
t.Fatalf("expected zero value, got %q", val)
76+
}
77+
}
78+
79+
func TestOptionalParam_TypeMismatch(t *testing.T) {
80+
r := newRequest(map[string]any{"n": 123})
81+
_, err := OptionalParam[string](r, "n")
82+
if err == nil {
83+
t.Fatalf("expected type mismatch error, got nil")
84+
}
85+
}
86+
87+
func TestOptionalParam_Success(t *testing.T) {
88+
r := newRequest(map[string]any{"s": "ok"})
89+
val, err := OptionalParam[string](r, "s")
90+
if err != nil {
91+
t.Fatalf("unexpected err: %v", err)
92+
}
93+
if val != "ok" {
94+
t.Fatalf("expected 'ok', got %q", val)
95+
}
96+
}
Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
package middleware
2+
3+
import (
4+
"context"
5+
6+
"github.com/harness/harness-mcp/cmd/harness-mcp-server/config"
7+
"github.com/harness/harness-mcp/pkg/harness/common"
8+
"github.com/mark3labs/mcp-go/mcp"
9+
"github.com/mark3labs/mcp-go/server"
10+
)
11+
12+
// WithHarnessScope creates a middleware that extracts the scope from the request
13+
// and adds it to the context for all tool handlers
14+
func WithHarnessScope(config *config.Config) server.ToolHandlerMiddleware {
15+
return func(next server.ToolHandlerFunc) server.ToolHandlerFunc {
16+
return func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
17+
// Extract scope from the request
18+
scope, err := common.FetchScope(config, request, false)
19+
if err != nil {
20+
// If we can't get the scope, continue with the request
21+
// The tool handler will handle this error if needed
22+
return next(ctx, request)
23+
}
24+
25+
// Add the entire scope to the context
26+
ctx = common.WithScopeContext(ctx, scope)
27+
28+
// Call the next handler with the enriched context
29+
return next(ctx, request)
30+
}
31+
}
32+
}

0 commit comments

Comments
 (0)