Re-read user session and refresh token when using MCP (#356)

Author: tomasmik

internal/cmd/audittrail/list.go

@@ -182,7 +182,7 @@ func searchAuditTrailEntries(ctx context.Context, input structs.SearchInput) (se
 		} `graphql:"searchAuditTrailEntries(input: $input)"`
 	}
 
-	if err := authenticated.Client.Query(
+	if err := authenticated.Client().Query(
 		ctx,
 		&query,
 		map[string]interface{}{"input": input},

internal/cmd/authenticated/client.go

@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"sync"
 
 	"github.com/urfave/cli/v3"
 
@@ -33,11 +34,30 @@ var (
 )
 
 // Client is the authenticated client that can be used by all CLI commands.
-var Client client.Client
+var (
+	auth client.Client
+	m    sync.Mutex
+)
+
+// Client returns the authenticated client.
+//
+// This is an unfortunate global which we have to lock for MCP.
+// TODO: Refactor this to not use a global.
+func Client() client.Client {
+	m.Lock()
+	defer m.Unlock()
+
+	return auth
+}
 
 // Ensure is a way of ensuring that the Client exists, and it meant to be used
 // as a Before action for commands that need it.
+//
+// You can also use it diretly to refresh the client.
 func Ensure(ctx context.Context, _ *cli.Command) (context.Context, error) {
+	m.Lock()
+	defer m.Unlock()
+
 	httpClient := client.GetHTTPClient()
 
 	if err := configureTLS(httpClient); err != nil {
@@ -49,7 +69,7 @@ func Ensure(ctx context.Context, _ *cli.Command) (context.Context, error) {
 		return ctx, err
 	}
 
-	Client = client.New(httpClient, session)
+	auth = client.New(httpClient, session)
 
 	return ctx, nil
 }

internal/cmd/authenticated/viewer.go

@@ -17,7 +17,7 @@ func CurrentViewer(ctx context.Context) (*Viewer, error) {
 	var query struct {
 		Viewer *Viewer
 	}
-	if err := Client.Query(ctx, &query, map[string]interface{}{}); err != nil {
+	if err := Client().Query(ctx, &query, map[string]interface{}{}); err != nil {
 		return nil, errors.Wrap(err, "failed to query user information")
 	}
 	if query.Viewer == nil {

internal/cmd/blueprint/deploy.go

@@ -77,7 +77,7 @@ func (c *deployCommand) deploy(ctx context.Context, cliCmd *cli.Command) error {
 		} `graphql:"blueprintCreateStack(id: $id, input: $input)"`
 	}
 
-	err = authenticated.Client.Mutate(
+	err = authenticated.Client().Mutate(
 		ctx,
 		&mutation,
 		map[string]any{
@@ -91,7 +91,7 @@ func (c *deployCommand) deploy(ctx context.Context, cliCmd *cli.Command) error {
 		return fmt.Errorf("failed to deploy stack from the blueprint: %w", err)
 	}
 
-	url := authenticated.Client.URL("/stack/%s", mutation.BlueprintCreateStack.StackID)
+	url := authenticated.Client().URL("/stack/%s", mutation.BlueprintCreateStack.StackID)
 	fmt.Printf("\nCreated stack: %q", url)
 
 	return nil

internal/cmd/blueprint/list.go

@@ -207,7 +207,7 @@ func searchBlueprints(ctx context.Context, input structs.SearchInput) (searchBlu
 		} `graphql:"searchBlueprints(input: $input)"`
 	}
 
-	if err := authenticated.Client.Query(
+	if err := authenticated.Client().Query(
 		ctx,
 		&query,
 		map[string]interface{}{"input": input},

internal/cmd/blueprint/show.go

@@ -164,7 +164,7 @@ func getBlueprintByID(ctx context.Context, blueprintID string) (blueprint, bool,
 		"blueprintId": graphql.ID(blueprintID),
 	}
 
-	if err := authenticated.Client.Query(ctx, &query, variables); err != nil {
+	if err := authenticated.Client().Query(ctx, &query, variables); err != nil {
 		return blueprint{}, false, errors.Wrapf(err, "failed to query for blueprint ID %q", blueprintID)
 	}
 

internal/cmd/draw/data/workerpools.go

@@ -21,7 +21,7 @@ type WorkerPool struct {
 
 // Selected opens the selected worker pool in the browser.
 func (q *WorkerPool) Selected(row table.Row) error {
-	return browser.OpenURL(authenticated.Client.URL("/stack/%s/run/%s", row[1], row[2]))
+	return browser.OpenURL(authenticated.Client().URL("/stack/%s/run/%s", row[1], row[2]))
 }
 
 // Columns returns the columns of the worker pool table.
@@ -73,7 +73,7 @@ func (q *WorkerPool) getPublicPoolRuns(ctx context.Context) ([]runsEdge, error)
 		} `graphql:"publicWorkerPool"`
 	}
 
-	if err := authenticated.Client.Query(ctx, &query, q.baseSearchParams()); err != nil {
+	if err := authenticated.Client().Query(ctx, &query, q.baseSearchParams()); err != nil {
 		return nil, errors.Wrap(err, "failed to query run list")
 	}
 
@@ -90,7 +90,7 @@ func (q *WorkerPool) getPrivatePoolRuns(ctx context.Context) ([]runsEdge, error)
 	vars := q.baseSearchParams()
 	vars["id"] = q.WokerPoolID
 
-	if err := authenticated.Client.Query(ctx, &query, vars); err != nil {
+	if err := authenticated.Client().Query(ctx, &query, vars); err != nil {
 		return nil, errors.Wrap(err, "failed to query run list")
 	}
 

internal/cmd/graphql/mcp.go

@@ -38,6 +38,7 @@ func registerIntrospectSchemaTool(s *server.MCPServer) {
 	)
 
 	s.AddTool(introspectTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		authenticated.Ensure(ctx, nil)
 		format := request.GetString("format", "summary")
 
 		var query struct {
@@ -120,7 +121,7 @@ func registerIntrospectSchemaTool(s *server.MCPServer) {
 			} `graphql:"__schema"`
 		}
 
-		if err := authenticated.Client.Query(ctx, &query, map[string]any{}); err != nil {
+		if err := authenticated.Client().Query(ctx, &query, map[string]any{}); err != nil {
 			return nil, errors.Wrap(err, "failed to introspect GraphQL schema")
 		}
 
@@ -144,6 +145,7 @@ func registerGetTypeDetailsTool(s *server.MCPServer) {
 	)
 
 	s.AddTool(typeDetailsTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		authenticated.Ensure(ctx, nil)
 		typeName, err := request.RequireString("type_name")
 		if err != nil {
 			return nil, err
@@ -191,7 +193,7 @@ func registerGetTypeDetailsTool(s *server.MCPServer) {
 			} `graphql:"__type(name: $name)"`
 		}
 
-		if err := authenticated.Client.Query(ctx, &query, map[string]any{"name": graphql.String(typeName)}); err != nil {
+		if err := authenticated.Client().Query(ctx, &query, map[string]any{"name": graphql.String(typeName)}); err != nil {
 			return nil, errors.Wrap(err, "failed to get type details")
 		}
 
@@ -223,6 +225,7 @@ func registerSearchSchemaFieldsTool(s *server.MCPServer) {
 	)
 
 	s.AddTool(searchTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		authenticated.Ensure(ctx, nil)
 		searchTerm, err := request.RequireString("search_term")
 		if err != nil {
 			return nil, err
@@ -271,7 +274,7 @@ func registerSearchSchemaFieldsTool(s *server.MCPServer) {
 			} `graphql:"__schema"`
 		}
 
-		if err := authenticated.Client.Query(ctx, &query, map[string]any{}); err != nil {
+		if err := authenticated.Client().Query(ctx, &query, map[string]any{}); err != nil {
 			return nil, errors.Wrap(err, "failed to introspect GraphQL schema")
 		}
 
@@ -505,6 +508,7 @@ func registerAuthenticationGuideTool(s *server.MCPServer) {
 	)
 
 	s.AddTool(authTool, func(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
+		authenticated.Ensure(ctx, nil)
 		authMethod := request.GetString("auth_method", "all")
 
 		var guide strings.Builder

internal/cmd/module/create_version.go

@@ -37,7 +37,7 @@ func createVersion(ctx context.Context, cliCmd *cli.Command) error {
 		"version":   version,
 	}
 
-	if err := authenticated.Client.Mutate(ctx, &mutation, variables); err != nil {
+	if err := authenticated.Client().Mutate(ctx, &mutation, variables); err != nil {
 		return err
 	}
 

internal/cmd/module/delete_version.go

@@ -25,7 +25,7 @@ func deleteVersion(ctx context.Context, cliCmd *cli.Command) error {
 		"module": graphql.ID(moduleID),
 	}
 
-	if err := authenticated.Client.Mutate(ctx, &mutation, variables); err != nil {
+	if err := authenticated.Client().Mutate(ctx, &mutation, variables); err != nil {
 		return err
 	}