Skip to content

Commit 1851665

Browse files
taskbotCopilot
andcommitted
parent fb475a8
author taskbot <taskbot@users.noreply.github.com> 1766072123 +0100 committer taskbot <taskbot@users.noreply.github.com> 1766158585 +0100 Integrate health monitoring into vMCP server Integrates the health monitoring infrastructure (from previous into the vMCP server, enabling periodic backend health checks with configurable Related-to: #3036 intervals and thresholds. changes from review changes from review add missing method Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
1 parent c5e4c3f commit 1851665

7 files changed

Lines changed: 662 additions & 26 deletions

File tree

cmd/vmcp/app/commands.go

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
vmcpclient "github.com/stacklok/toolhive/pkg/vmcp/client"
2121
"github.com/stacklok/toolhive/pkg/vmcp/config"
2222
"github.com/stacklok/toolhive/pkg/vmcp/discovery"
23+
"github.com/stacklok/toolhive/pkg/vmcp/health"
2324
vmcprouter "github.com/stacklok/toolhive/pkg/vmcp/router"
2425
vmcpserver "github.com/stacklok/toolhive/pkg/vmcp/server"
2526
)
@@ -253,6 +254,8 @@ func discoverBackends(ctx context.Context, cfg *config.Config) ([]vmcp.Backend,
253254
}
254255

255256
// runServe implements the serve command logic
257+
//
258+
//nolint:gocyclo // Complexity from server initialization and configuration is acceptable
256259
func runServe(cmd *cobra.Command, _ []string) error {
257260
ctx := cmd.Context()
258261
configPath := viper.GetString("config")
@@ -330,15 +333,29 @@ func runServe(cmd *cobra.Command, _ []string) error {
330333
}()
331334
}
332335

336+
// Configure health monitoring if enabled
337+
var healthMonitorConfig *health.MonitorConfig
338+
if cfg.Operational != nil && cfg.Operational.FailureHandling != nil && cfg.Operational.FailureHandling.HealthCheckInterval > 0 {
339+
defaults := health.DefaultConfig()
340+
healthMonitorConfig = &health.MonitorConfig{
341+
CheckInterval: time.Duration(cfg.Operational.FailureHandling.HealthCheckInterval),
342+
UnhealthyThreshold: cfg.Operational.FailureHandling.UnhealthyThreshold,
343+
Timeout: defaults.Timeout,
344+
DegradedThreshold: defaults.DegradedThreshold,
345+
}
346+
logger.Info("Health monitoring configured from operational settings")
347+
}
348+
333349
serverCfg := &vmcpserver.Config{
334-
Name: cfg.Name,
335-
Version: getVersion(),
336-
Host: host,
337-
Port: port,
338-
AuthMiddleware: authMiddleware,
339-
AuthInfoHandler: authInfoHandler,
340-
TelemetryProvider: telemetryProvider,
341-
AuditConfig: cfg.Audit,
350+
Name: cfg.Name,
351+
Version: getVersion(),
352+
Host: host,
353+
Port: port,
354+
AuthMiddleware: authMiddleware,
355+
AuthInfoHandler: authInfoHandler,
356+
TelemetryProvider: telemetryProvider,
357+
AuditConfig: cfg.Audit,
358+
HealthMonitorConfig: healthMonitorConfig,
342359
}
343360

344361
// Convert composite tool configurations to workflow definitions

pkg/vmcp/auth/strategies/header_injection.go

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88

99
"github.com/stacklok/toolhive/pkg/validation"
1010
authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types"
11+
"github.com/stacklok/toolhive/pkg/vmcp/health"
1112
)
1213

1314
// HeaderInjectionStrategy injects a static header value into request headers.
@@ -46,20 +47,26 @@ func (*HeaderInjectionStrategy) Name() string {
4647
// Authenticate injects the header value from the strategy config into the request header.
4748
//
4849
// This method:
49-
// 1. Validates that HeaderName and HeaderValue are present in the strategy config
50-
// 2. Sets the specified header with the provided value
50+
// 1. Skips authentication if this is a health check request
51+
// 2. Validates that HeaderName and HeaderValue are present in the strategy config
52+
// 3. Sets the specified header with the provided value
5153
//
5254
// Parameters:
53-
// - ctx: Request context (currently unused, reserved for future secret resolution)
55+
// - ctx: Request context (used to check for health check marker)
5456
// - req: The HTTP request to authenticate
5557
// - strategy: The backend auth strategy configuration containing HeaderInjection
5658
//
5759
// Returns an error if:
5860
// - HeaderName is missing or empty
5961
// - HeaderValue is missing or empty
6062
func (*HeaderInjectionStrategy) Authenticate(
61-
_ context.Context, req *http.Request, strategy *authtypes.BackendAuthStrategy,
63+
ctx context.Context, req *http.Request, strategy *authtypes.BackendAuthStrategy,
6264
) error {
65+
// Skip authentication for health checks
66+
if health.IsHealthCheck(ctx) {
67+
return nil
68+
}
69+
6370
if strategy == nil || strategy.HeaderInjection == nil {
6471
return fmt.Errorf("header_injection configuration required")
6572
}

pkg/vmcp/auth/strategies/tokenexchange.go

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import (
1212
"github.com/stacklok/toolhive/pkg/auth/tokenexchange"
1313
"github.com/stacklok/toolhive/pkg/env"
1414
authtypes "github.com/stacklok/toolhive/pkg/vmcp/auth/types"
15+
"github.com/stacklok/toolhive/pkg/vmcp/health"
1516
)
1617

1718
const (
@@ -69,18 +70,19 @@ func (*TokenExchangeStrategy) Name() string {
6970
// Authenticate exchanges the client's token for a backend token and injects it.
7071
//
7172
// This method:
72-
// 1. Retrieves the client's identity and token from the context
73-
// 2. Parses and validates the token exchange configuration from strategy
74-
// 3. Gets or creates a cached ExchangeConfig for this backend configuration
75-
// 4. Creates a TokenSource with the current identity token
76-
// 5. Obtains an access token by performing the exchange
77-
// 6. Injects the token into the backend request's Authorization header
73+
// 1. Skips authentication if this is a health check request
74+
// 2. Retrieves the client's identity and token from the context
75+
// 3. Parses and validates the token exchange configuration from strategy
76+
// 4. Gets or creates a cached ExchangeConfig for this backend configuration
77+
// 5. Creates a TokenSource with the current identity token
78+
// 6. Obtains an access token by performing the exchange
79+
// 7. Injects the token into the backend request's Authorization header
7880
//
7981
// Token caching per user is handled by the upper vMCP TokenCache layer.
8082
// This strategy only caches the ExchangeConfig template per backend.
8183
//
8284
// Parameters:
83-
// - ctx: Request context containing the authenticated identity
85+
// - ctx: Request context containing the authenticated identity (or health check marker)
8486
// - req: The HTTP request to authenticate
8587
// - strategy: Backend auth strategy containing token exchange configuration
8688
//
@@ -92,6 +94,11 @@ func (*TokenExchangeStrategy) Name() string {
9294
func (s *TokenExchangeStrategy) Authenticate(
9395
ctx context.Context, req *http.Request, strategy *authtypes.BackendAuthStrategy,
9496
) error {
97+
// Skip authentication for health checks
98+
if health.IsHealthCheck(ctx) {
99+
return nil
100+
}
101+
95102
identity, ok := auth.IdentityFromContext(ctx)
96103
if !ok {
97104
return fmt.Errorf("no identity found in context")

pkg/vmcp/health/monitor.go

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,24 @@ import (
1010
"github.com/stacklok/toolhive/pkg/vmcp"
1111
)
1212

13+
// healthCheckContextKey is a marker for health check requests.
14+
type healthCheckContextKey struct{}
15+
16+
// WithHealthCheckMarker marks a context as a health check request.
17+
// Authentication layers can use IsHealthCheck to identify and skip authentication
18+
// for health check requests.
19+
func WithHealthCheckMarker(ctx context.Context) context.Context {
20+
return context.WithValue(ctx, healthCheckContextKey{}, true)
21+
}
22+
23+
// IsHealthCheck returns true if the context is marked as a health check.
24+
// Authentication strategies use this to bypass authentication for health checks,
25+
// since health checks verify backend availability and should not require user credentials.
26+
func IsHealthCheck(ctx context.Context) bool {
27+
val, ok := ctx.Value(healthCheckContextKey{}).(bool)
28+
return ok && val
29+
}
30+
1331
// Monitor performs periodic health checks on backend MCP servers.
1432
// It runs background goroutines for each backend, tracking their health status
1533
// and consecutive failure counts. The monitor supports graceful shutdown and
@@ -219,8 +237,12 @@ func (m *Monitor) performHealthCheck(ctx context.Context, backend *vmcp.Backend)
219237
Metadata: backend.Metadata,
220238
}
221239

240+
// Mark context as health check to bypass authentication
241+
// Health checks verify backend availability and should not require user credentials
242+
healthCheckCtx := WithHealthCheckMarker(ctx)
243+
222244
// Perform health check
223-
status, err := m.checker.CheckHealth(ctx, target)
245+
status, err := m.checker.CheckHealth(healthCheckCtx, target)
224246

225247
// Record result in status tracker
226248
if err != nil {

0 commit comments

Comments
 (0)