chore: fix TestManagedAgentLimit flake (#19026)

Closes https://github.com/coder/internal/issues/812

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)
 				})
 			}
 

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)