From 9a05b4679b9bc4898d5f9fb2e6089957976fb27b Mon Sep 17 00:00:00 2001
From: Dean Sheather <dean@deansheather.com>
Date: Thu, 24 Jul 2025 15:13:15 +1000
Subject: [PATCH] chore: fix TestManagedAgentLimit flake (#19026)

Closes https://github.com/coder/internal/issues/812
---
 enterprise/coderd/coderd.go          | 4 ++--
 enterprise/coderd/license/license.go | 9 +++++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
index 16ab9c77c7..9583e14cd7 100644
--- a/enterprise/coderd/coderd.go
+++ b/enterprise/coderd/coderd.go
@@ -830,7 +830,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 						}
 						api.derpMesh.SetAddresses(addresses, false)
 					}
-					_ = api.updateEntitlements(ctx)
+					_ = api.updateEntitlements(api.ctx)
 				})
 			} else {
 				coordinator = agpltailnet.NewCoordinator(api.Logger)
@@ -840,7 +840,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 				api.replicaManager.SetCallback(func() {
 					// If the amount of replicas change, so should our entitlements.
 					// This is to display a warning in the UI if the user is unlicensed.
-					_ = api.updateEntitlements(ctx)
+					_ = api.updateEntitlements(api.ctx)
 				})
 			}
 
diff --git a/enterprise/coderd/license/license.go b/enterprise/coderd/license/license.go
index 6b31daa72a..bc5c174d9f 100644
--- a/enterprise/coderd/license/license.go
+++ b/enterprise/coderd/license/license.go
@@ -432,10 +432,15 @@ func LicensesEntitlements(
 		if featureArguments.ManagedAgentCountFn != nil {
 			managedAgentCount, err = featureArguments.ManagedAgentCountFn(ctx, agentLimit.UsagePeriod.Start, agentLimit.UsagePeriod.End)
 		}
-		if err != nil {
+		switch {
+		case xerrors.Is(err, context.Canceled) || xerrors.Is(err, context.DeadlineExceeded):
+			// If the context is canceled, we want to bail the entire
+			// LicensesEntitlements call.
+			return entitlements, xerrors.Errorf("get managed agent count: %w", err)
+		case err != nil:
 			entitlements.Errors = append(entitlements.Errors,
 				fmt.Sprintf("Error getting managed agent count: %s", err.Error()))
-		} else {
+		default:
 			agentLimit.Actual = &managedAgentCount
 			entitlements.AddFeature(codersdk.FeatureManagedAgentLimit, agentLimit)