fix(aws-marketplace): check the container license with CheckoutLicense so AWS's standard buyer policy works (#9981)

* fix(aws-marketplace): check the container license with CheckoutLicense so AWS's standard buyer policy works

A buyer who follows AWS's generated launch steps gets the managed
AWSLicenseManagerConsumptionPolicy, which does not allow
ListReceivedLicenses, so the check failed and Sim stayed on the open-source
core. Check out a provisional unit of the entitlement and return it at once,
the integration AWS documents for contract-priced containers.

License Manager answers a unit held by another replica, or one a crashed
replica never returned, the same way as a missing entitlement. Startup retries
that answer, and a running Enterprise deployment only restarts without
Enterprise once it has held longer than an unreturned unit can last.

* fix(aws-marketplace): count only an unbroken run of no-entitlement answers, assert the checkout request, bump the chart

* fix(aws-marketplace): check the container license as an uncounted tier entitlement

AWS documents a contract that unlocks a tier of features, rather than metering
a quantity, as a provisional CheckoutLicense with unit None and no value. That
checkout is not counted against the contract, so replicas never compete for a
unit, and the check-in, startup retry, and one-hour confirmation window that
guarded against a held unit go away. A definite answer that differs from the
boot mode restarts; a failed check keeps the current mode.

Author: waleedlatif1

apps/sim/lib/billing/aws-marketplace/container-license.test.ts

@@ -0,0 +1,124 @@
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+
+const { mockSend } = vi.hoisted(() => ({ mockSend: vi.fn() }))
+
+vi.mock('@aws-sdk/client-license-manager', () => ({
+  LicenseManagerClient: class {
+    send = mockSend
+  },
+  CheckoutLicenseCommand: class {
+    constructor(public input: { ProductSKU: string }) {}
+  },
+}))
+vi.mock('@/lib/core/network/proxy-agent.server', () => ({ awsRequestHandler: () => undefined }))
+vi.mock('@sim/utils/random', () => ({ randomInt: (min: number) => min }))
+
+import { applyAwsMarketplaceContainerLicense } from '@/lib/billing/aws-marketplace/container-license'
+
+const MINUTE = 60 * 1000
+
+type Answer =
+  | 'entitled'
+  | 'NoEntitlementsAllowedException'
+  | 'ResourceNotFoundException'
+  | 'AccessDeniedException'
+  | 'ServerInternalException'
+
+/** License Manager's checkout answers in order; the last one repeats once the list runs out. */
+let answers: Answer[] = []
+const checkoutRequests: unknown[] = []
+
+function awsError(name: string): Error {
+  return Object.assign(new Error(name), { name })
+}
+
+function enterpriseEnabled(): boolean {
+  return process.env.ENTERPRISE_ENABLED === 'true'
+}
+
+describe('applyAwsMarketplaceContainerLicense', () => {
+  let kill: ReturnType<typeof vi.spyOn>
+
+  beforeEach(() => {
+    vi.useFakeTimers({ toFake: ['setTimeout'] })
+    answers = []
+    checkoutRequests.length = 0
+    vi.stubEnv('AWS_MARKETPLACE_LICENSE_PRODUCT_SKU', 'product-placeholder')
+    kill = vi.spyOn(process, 'kill').mockImplementation(() => true)
+    mockSend.mockReset()
+    mockSend.mockImplementation(async (command: { input: unknown }) => {
+      checkoutRequests.push(command.input)
+      const answer = answers.length > 1 ? answers.shift() : answers[0]
+      if (answer !== 'entitled') throw awsError(answer ?? 'ServerInternalException')
+      return { EntitlementsAllowed: [{ Name: 'enterprise_license', Unit: 'None' }] }
+    })
+  })
+
+  afterEach(() => {
+    vi.clearAllTimers()
+    vi.useRealTimers()
+    kill.mockRestore()
+    Reflect.deleteProperty(process.env, 'ENTERPRISE_ENABLED')
+    Reflect.deleteProperty(process.env, 'NEXT_PUBLIC_ENTERPRISE_ENABLED')
+  })
+
+  it('does nothing outside a Marketplace install', async () => {
+    vi.stubEnv('AWS_MARKETPLACE_LICENSE_PRODUCT_SKU', undefined)
+    await applyAwsMarketplaceContainerLicense()
+    expect(mockSend).not.toHaveBeenCalled()
+  })
+
+  it('turns Enterprise on with an uncounted checkout of the entitlement', async () => {
+    answers = ['entitled']
+    await applyAwsMarketplaceContainerLicense()
+    expect(enterpriseEnabled()).toBe(true)
+    expect(checkoutRequests).toEqual([
+      expect.objectContaining({
+        ProductSKU: 'product-placeholder',
+        CheckoutType: 'PROVISIONAL',
+        KeyFingerprint: 'aws:294406891311:AWS/Marketplace:issuer-fingerprint',
+        Entitlements: [{ Name: 'enterprise_license', Unit: 'None' }],
+      }),
+    ])
+  })
+
+  it('runs without Enterprise when no license grants the entitlement', async () => {
+    answers = ['NoEntitlementsAllowedException']
+    await applyAwsMarketplaceContainerLicense()
+    expect(enterpriseEnabled()).toBe(false)
+  })
+
+  it('runs without Enterprise when its IAM role cannot check out', async () => {
+    answers = ['AccessDeniedException']
+    await applyAwsMarketplaceContainerLicense()
+    expect(enterpriseEnabled()).toBe(false)
+  })
+
+  it.each<Answer>(['NoEntitlementsAllowedException', 'ResourceNotFoundException'])(
+    'restarts without Enterprise once the license ends (%s)',
+    async (answer) => {
+      answers = ['entitled', answer]
+      await applyAwsMarketplaceContainerLicense()
+      await vi.advanceTimersByTimeAsync(16 * MINUTE)
+      expect(kill).toHaveBeenCalledWith(process.pid, 'SIGTERM')
+    }
+  )
+
+  it('restarts into Enterprise once the license becomes active', async () => {
+    answers = ['NoEntitlementsAllowedException', 'entitled']
+    await applyAwsMarketplaceContainerLicense()
+    expect(enterpriseEnabled()).toBe(false)
+    await vi.advanceTimersByTimeAsync(16 * MINUTE)
+    expect(kill).toHaveBeenCalledWith(process.pid, 'SIGTERM')
+  })
+
+  it.each<Answer>(['ServerInternalException', 'AccessDeniedException'])(
+    'never restarts a running Enterprise deployment over a failed check (%s)',
+    async (answer) => {
+      answers = ['entitled', answer]
+      await applyAwsMarketplaceContainerLicense()
+      await vi.advanceTimersByTimeAsync(6 * 60 * MINUTE)
+      expect(kill).not.toHaveBeenCalled()
+    }
+  )
+})

apps/sim/lib/billing/aws-marketplace/container-license.ts

@@ -1,26 +1,21 @@
-import {
-  type GrantedLicense,
-  LicenseManagerClient,
-  ListReceivedLicensesCommand,
-} from '@aws-sdk/client-license-manager'
+import { CheckoutLicenseCommand, LicenseManagerClient } from '@aws-sdk/client-license-manager'
 import { createLogger } from '@sim/logger'
 import { getErrorMessage } from '@sim/utils/errors'
+import { generateId } from '@sim/utils/id'
 import { randomInt } from '@sim/utils/random'
 import { awsRequestHandler } from '@/lib/core/network/proxy-agent.server'
 
 /**
  * Enterprise entitlement for Sim installed from the AWS Marketplace container listing.
  *
  * The listing sells an annual contract whose one dimension is `enterprise_license`. AWS Marketplace
- * grants the buyer a license in AWS License Manager, and the running product reads it. While the
- * license is active, the Enterprise suite is on (`ENTERPRISE_ENABLED`). Without one, Sim runs as the
- * open-source core, so a lapsed renewal or an unreachable License Manager never locks a buyer out
- * of their deployment.
- *
- * The check reads the license rather than checking a unit out of it. A contract usually covers the
- * whole deployment with one unit, so checking it out from every replica would make them compete
- * for it, and a unit left checked out (a failed check-in holds it for an hour) would read exactly
- * like an expired license.
+ * grants the buyer a license in AWS License Manager, and the running product checks it the way AWS
+ * documents for a contract that unlocks a tier of features rather than metering a quantity: a
+ * provisional `CheckoutLicense` of the entitlement with unit `None`. Such a checkout is not counted
+ * against the contract, so every replica can check at once, and it needs only the permission AWS's
+ * managed `AWSLicenseManagerConsumptionPolicy` grants. While the license is active, the Enterprise
+ * suite is on (`ENTERPRISE_ENABLED`). Without one, Sim runs as the open-source core, so a lapsed
+ * renewal or an unreachable License Manager never locks a buyer out of their deployment.
  *
  * Opt-in: nothing here runs unless {@link PRODUCT_SKU_ENV} is set, which only the Marketplace
  * Helm profile does. Enterprise flags are read once at startup by both the server and the
@@ -34,8 +29,8 @@ const logger = createLogger('AwsMarketplaceContainerLicense')
 const PRODUCT_SKU_ENV = 'AWS_MARKETPLACE_LICENSE_PRODUCT_SKU'
 
 /**
- * Optional override for where the license is read. License Manager replicates a granted license to
- * other Regions, so the default is the pod's own Region, the same chain AWS's EKS walkthrough uses.
+ * Optional override for where the license is checked. License Manager replicates a granted license
+ * to other Regions, so the default is the pod's own Region, the same chain AWS's EKS walkthrough uses.
  */
 const LICENSE_REGION_ENV = 'AWS_MARKETPLACE_LICENSE_REGION'
 
@@ -48,21 +43,26 @@ const ENTERPRISE_ENTITLEMENT = 'enterprise_license'
 /** AWS recommends re-checking every 15 minutes to pick up renewals, upgrades, and cancellations. */
 const RECHECK_INTERVAL_MS = 15 * 60 * 1000
 
-/**
- * Replicas start together, so their re-checks line up; each waits a random part of this before
- * restarting, so a renewal or lapse never takes every replica down at once.
- */
-const RESTART_SPREAD_MS = 10 * 60 * 1000
+/** Replicas start together; a random part of this added to each re-check keeps them apart. */
+const RECHECK_JITTER_MS = 2 * 60 * 1000
 
 /**
- * How the buyer's account holds the license. A delegate grant only shows licenses distributed to
- * other accounts, and a grant still awaiting acceptance, or rejected, disabled, or deleted by the
- * license administrator, entitles nothing. No status at all is a license the account bought itself.
+ * Replicas see a change together; each waits a random part of this before restarting, so a renewal
+ * or lapse never takes every replica down at once.
  */
-const USABLE_RECEIVED_STATUSES = new Set(['ACTIVE', 'WORKFLOW_COMPLETED'])
+const RESTART_SPREAD_MS = 10 * 60 * 1000
 
 const ENTERPRISE_FLAGS = ['ENTERPRISE_ENABLED', 'NEXT_PUBLIC_ENTERPRISE_ENABLED'] as const
 
+/** License Manager's answers when the account holds no active license with the entitlement. */
+const NO_ENTITLEMENT_ERRORS = new Set([
+  'NoEntitlementsAllowedException',
+  'ResourceNotFoundException',
+])
+
+/** The ServiceAccount's IAM role does not allow the check: a setup problem, never a verdict. */
+const PERMISSION_ERRORS = new Set(['AccessDeniedException', 'AuthorizationException'])
+
 /** `unknown` is a failed check (network, permissions, throttling), never a verdict on the license. */
 type LicenseState = 'entitled' | 'not-entitled' | 'unknown'
 
@@ -78,58 +78,62 @@ function createClient(): LicenseManagerClient {
   })
 }
 
-function isUsableGrant({ ReceivedMetadata: received }: GrantedLicense): boolean {
-  if (received?.GrantType === 'DelegateGrant') return false
-  return !received?.ReceivedStatus || USABLE_RECEIVED_STATUSES.has(received.ReceivedStatus)
-}
-
-/**
- * An AWS Marketplace license this product's buyer holds is current when License Manager reports it
- * `AVAILABLE`, its issuer is AWS Marketplace, it carries the Enterprise entitlement, and the account
- * holds it through a grant it may use. License Manager moves an ended, cancelled, or suspended
- * license out of `AVAILABLE` itself.
- */
-function isActiveEnterpriseLicense(license: GrantedLicense): boolean {
-  return (
-    license.Status === 'AVAILABLE' &&
-    isUsableGrant(license) &&
-    license.Issuer?.KeyFingerprint === MARKETPLACE_KEY_FINGERPRINT &&
-    (license.Entitlements ?? []).some((entitlement) => entitlement.Name === ENTERPRISE_ENTITLEMENT)
-  )
+function classifyFailure(error: unknown): LicenseState {
+  const name = error instanceof Error ? error.name : undefined
+  if (name && NO_ENTITLEMENT_ERRORS.has(name)) {
+    logger.info('AWS License Manager reports no active Enterprise entitlement', { error: name })
+    return 'not-entitled'
+  }
+  if (name && PERMISSION_ERRORS.has(name)) {
+    logger.error(
+      'The IAM role cannot check the AWS Marketplace license; allow license-manager:CheckoutLicense',
+      { error: name }
+    )
+    return 'unknown'
+  }
+  logger.warn('Could not check the AWS Marketplace Enterprise license', {
+    error: name,
+    message: getErrorMessage(error, 'Unknown error'),
+  })
+  return 'unknown'
 }
 
 async function checkLicense(
   client: LicenseManagerClient,
   productSku: string
 ): Promise<LicenseState> {
   try {
-    let nextToken: string | undefined
-    do {
-      const page = await client.send(
-        new ListReceivedLicensesCommand({
-          Filters: [
-            { Name: 'ProductSKU', Values: [productSku] },
-            { Name: 'Fingerprint', Values: [MARKETPLACE_KEY_FINGERPRINT] },
-          ],
-          NextToken: nextToken,
-        })
-      )
-      if ((page.Licenses ?? []).some(isActiveEnterpriseLicense)) return 'entitled'
-      nextToken = page.NextToken
-    } while (nextToken)
-    return 'not-entitled'
+    const checkout = await client.send(
+      new CheckoutLicenseCommand({
+        ProductSKU: productSku,
+        CheckoutType: 'PROVISIONAL',
+        KeyFingerprint: MARKETPLACE_KEY_FINGERPRINT,
+        Entitlements: [{ Name: ENTERPRISE_ENTITLEMENT, Unit: 'None' }],
+        ClientToken: generateId(),
+      })
+    )
+    const allowed = (checkout.EntitlementsAllowed ?? []).some(
+      (entitlement) => entitlement.Name === ENTERPRISE_ENTITLEMENT
+    )
+    return allowed ? 'entitled' : 'not-entitled'
   } catch (error) {
-    logger.warn('Could not check the AWS Marketplace Enterprise license', {
-      message: getErrorMessage(error, 'Unknown error'),
-    })
-    return 'unknown'
+    return classifyFailure(error)
   }
 }
 
 function setEnterpriseFlags(enabled: boolean): void {
   for (const flag of ENTERPRISE_FLAGS) process.env[flag] = enabled ? 'true' : 'false'
 }
 
+function restartToApply(state: LicenseState): void {
+  const delayMs = randomInt(0, RESTART_SPREAD_MS)
+  logger.info('AWS Marketplace Enterprise license changed; restarting to apply it', {
+    state,
+    delayMs,
+  })
+  setTimeout(() => process.kill(process.pid, 'SIGTERM'), delayMs)
+}
+
 /**
  * Resolves the Enterprise license before the server loads, so every Enterprise gate reads the
  * answer through `ENTERPRISE_ENABLED`, then keeps re-checking. Returns without doing anything on a
@@ -150,17 +154,16 @@ export async function applyAwsMarketplaceContainerLicense(): Promise<void> {
     { state: bootState }
   )
 
-  const recheck = setInterval(async () => {
+  const scheduleRecheck = () => {
+    setTimeout(recheck, RECHECK_INTERVAL_MS + randomInt(0, RECHECK_JITTER_MS)).unref()
+  }
+  const recheck = async () => {
     const state = await checkLicense(client, productSku)
-    const entitledNow = state === 'entitled'
-    if (state === 'unknown' || entitledNow === enterpriseAtBoot) return
-    clearInterval(recheck)
-    const delayMs = randomInt(0, RESTART_SPREAD_MS)
-    logger.info('AWS Marketplace Enterprise license changed; restarting to apply it', {
-      state,
-      delayMs,
-    })
-    setTimeout(() => process.kill(process.pid, 'SIGTERM'), delayMs)
-  }, RECHECK_INTERVAL_MS)
-  recheck.unref()
+    // A failed check keeps the current mode; only a definite answer that differs restarts.
+    const changed =
+      (state === 'entitled' && !enterpriseAtBoot) || (state === 'not-entitled' && enterpriseAtBoot)
+    if (changed) return restartToApply(state)
+    scheduleRecheck()
+  }
+  scheduleRecheck()
 }

helm/sim/Chart.yaml

@@ -2,7 +2,7 @@ apiVersion: v2
 name: sim
 description: A Helm chart for Sim - the open-source AI workspace where teams build, deploy, and manage AI agents
 type: application
-version: 1.11.16
+version: 1.11.17
 appVersion: "v0.8.26"
 kubeVersion: ">=1.25.0-0"
 home: https://sim.ai

helm/sim/examples/aws-marketplace-prerequisites.yaml

@@ -159,9 +159,10 @@ Resources:
             Version: '2012-10-17'
             Statement:
               - Effect: Allow
-                # Sim only reads the license. ListReceivedLicenses has no resource-level
-                # permissions, so it cannot be scoped below '*'.
-                Action: license-manager:ListReceivedLicenses
+                # Sim checks the license with an uncounted checkout, which AWS's managed
+                # AWSLicenseManagerConsumptionPolicy also grants. CheckoutLicense has no
+                # resource-level permissions, so it cannot be scoped below '*'.
+                Action: license-manager:CheckoutLicense
                 Resource: '*'
 
 Outputs:

helm/sim/examples/values-aws-marketplace.yaml

@@ -21,8 +21,9 @@
 #      provider (for IRSA).
 #   2. RDS for PostgreSQL 17 and ElastiCache reachable from the cluster. The listing's
 #      CloudFormation template creates both in your VPC.
-#   3. An IAM role for the ServiceAccount below that allows license-manager:ListReceivedLicenses
-#      (the CloudFormation template creates it).
+#   3. An IAM role for the ServiceAccount below that allows license-manager:CheckoutLicense (the
+#      CloudFormation template creates it; AWS's managed AWSLicenseManagerConsumptionPolicy
+#      grants it too).
 #   4. An ACM certificate for the two hostnames below.
 #   5. Two Secrets in the install namespace (no secret ever goes in this file). The template's
 #      DatabasePasswordArn and CacheAuthTokenArn outputs are the ARNs of the Secrets Manager