diff --git a/internal/server/task_run_handlers.go b/internal/server/task_run_handlers.go index 2b332985..190fbe54 100644 --- a/internal/server/task_run_handlers.go +++ b/internal/server/task_run_handlers.go @@ -666,12 +666,17 @@ func (s *Server) buildBizdevTaskRequest(dto TaskRunRequestDTO, sessionID, client func resolveTaskBackend(s *Server, requested string) string { configured := "" + env := "" if s != nil && s.cfg != nil { configured = s.cfg.AgentBackend + env = s.cfg.Env } - resolved, err := taskrun.ResolveBackend(requested, configured) + resolved, err := taskrun.ResolveBackendInEnv(requested, configured, env) if err != nil { - fallback, _ := taskrun.ResolveBackend("", configured) + fallback, ferr := taskrun.ResolveBackendInEnv("", configured, env) + if ferr != nil { + return "" + } return fallback } return resolved diff --git a/internal/taskrun/backend.go b/internal/taskrun/backend.go index 3c95a975..67e4dbba 100644 --- a/internal/taskrun/backend.go +++ b/internal/taskrun/backend.go @@ -17,23 +17,50 @@ var allowedBackends = map[string]struct{}{ "cursor": {}, } +func isProductionEnv(env string) bool { + switch strings.ToLower(strings.TrimSpace(env)) { + case "main", "prod", "production": + return true + default: + return false + } +} + // ResolveBackend maps a caller hint onto the server allowlist. // An empty requested value falls back to configured. Unauthorized values are rejected. +// Prefer ResolveBackendInEnv when the deployment environment is known. func ResolveBackend(requested, configured string) (string, error) { + return ResolveBackendInEnv(requested, configured, "") +} + +// ResolveBackendInEnv is ResolveBackend with environment policy: +// in main/prod, stub is never allowed (neither as request nor as empty fallback). +func ResolveBackendInEnv(requested, configured, env string) (string, error) { requested = strings.TrimSpace(requested) configured = strings.TrimSpace(configured) + prod := isProductionEnv(env) + if requested == "" { if configured == "" { + if prod { + return "", ErrUnauthorizedBackend + } return "stub", nil } if _, ok := allowedBackends[configured]; !ok { return "", ErrUnauthorizedBackend } + if prod && configured == "stub" { + return "", ErrUnauthorizedBackend + } return configured, nil } if _, ok := allowedBackends[requested]; !ok { return "", ErrUnauthorizedBackend } + if prod && requested == "stub" { + return "", ErrUnauthorizedBackend + } return requested, nil } diff --git a/internal/taskrun/manifest_test.go b/internal/taskrun/manifest_test.go index a3657f22..f081a314 100644 --- a/internal/taskrun/manifest_test.go +++ b/internal/taskrun/manifest_test.go @@ -63,6 +63,26 @@ func TestResolveBackend_StripsUnauthorized(t *testing.T) { } } +func TestResolveBackendInEnv_RejectsStubInMain(t *testing.T) { + if _, err := ResolveBackendInEnv("stub", "pi-agent", "main"); err == nil { + t.Fatal("expected stub rejected in main") + } + if _, err := ResolveBackendInEnv("", "stub", "main"); err == nil { + t.Fatal("expected stub config rejected in main") + } + if _, err := ResolveBackendInEnv("", "", "main"); err == nil { + t.Fatal("expected empty fallback rejected in main") + } + got, err := ResolveBackendInEnv("stub", "pi-agent", "dev") + if err != nil || got != "stub" { + t.Fatalf("dev should allow stub, got %q %v", got, err) + } + got, err = ResolveBackendInEnv("", "pi-agent", "main") + if err != nil || got != "pi-agent" { + t.Fatalf("main should allow pi-agent, got %q %v", got, err) + } +} + func TestStripSecretsAndSensitive(t *testing.T) { meta := StripSecrets(map[string]any{"api_key": "secret", "run_id": "r1", "authorization": "Bearer x"}) if _, ok := meta["api_key"]; ok {