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

Author: waleedlatif1

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

@@ -22,12 +22,17 @@ import { applyAwsMarketplaceContainerLicense } from '@/lib/billing/aws-marketpla
 const MINUTE = 60 * 1000
 const TOKEN = 'consumption-token-placeholder'
 
-type Answer = 'entitled' | 'NoEntitlementsAllowedException' | 'AccessDeniedException'
+type Answer =
+  | 'entitled'
+  | 'NoEntitlementsAllowedException'
+  | 'AccessDeniedException'
+  | 'ServerInternalException'
 
 /** License Manager's checkout answers in order; the last one repeats once the list runs out. */
 let answers: Answer[] = []
 let checkInFails = false
-const checkIns: string[] = []
+const checkoutRequests: unknown[] = []
+const checkIns: Array<string | undefined> = []
 
 function awsError(name: string): Error {
   return Object.assign(new Error(name), { name })
@@ -51,6 +56,7 @@ describe('applyAwsMarketplaceContainerLicense', () => {
     answers = []
     checkInFails = false
     checkIns.length = 0
+    checkoutRequests.length = 0
     vi.stubEnv('AWS_MARKETPLACE_LICENSE_PRODUCT_SKU', 'product-placeholder')
     kill = vi.spyOn(process, 'kill').mockImplementation(() => true)
     mockSend.mockReset()
@@ -60,6 +66,7 @@ describe('applyAwsMarketplaceContainerLicense', () => {
         if (checkInFails) throw awsError('ServerInternalException')
         return {}
       }
+      checkoutRequests.push((command as { input: unknown }).input)
       const answer = answers.length > 1 ? answers.shift() : answers[0]
       if (answer !== 'entitled') throw awsError(answer ?? 'ServerInternalException')
       return {
@@ -73,7 +80,6 @@ describe('applyAwsMarketplaceContainerLicense', () => {
     vi.clearAllTimers()
     vi.useRealTimers()
     kill.mockRestore()
-    vi.unstubAllEnvs()
     Reflect.deleteProperty(process.env, 'ENTERPRISE_ENABLED')
     Reflect.deleteProperty(process.env, 'NEXT_PUBLIC_ENTERPRISE_ENABLED')
   })
@@ -88,6 +94,14 @@ describe('applyAwsMarketplaceContainerLicense', () => {
     answers = ['entitled']
     await boot()
     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', Value: '1', Unit: 'Count' }],
+      }),
+    ])
     expect(checkIns).toEqual([TOKEN])
   })
 
@@ -132,6 +146,19 @@ describe('applyAwsMarketplaceContainerLicense', () => {
     expect(kill).toHaveBeenCalledWith(process.pid, 'SIGTERM')
   })
 
+  it('does not count failed checks toward switching Enterprise off', async () => {
+    answers = [
+      'entitled',
+      'NoEntitlementsAllowedException',
+      ...Array<Answer>(6).fill('ServerInternalException'),
+      'NoEntitlementsAllowedException',
+      'entitled',
+    ]
+    await boot()
+    await vi.advanceTimersByTimeAsync(130 * MINUTE)
+    expect(kill).not.toHaveBeenCalled()
+  })
+
   it('never restarts a running Enterprise deployment over a permission failure', async () => {
     answers = ['entitled', 'AccessDeniedException']
     await boot()

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

@@ -61,8 +61,9 @@ const RECHECK_JITTER_MS = 2 * 60 * 1000
 const UNRETURNED_UNIT_MS = 60 * 60 * 1000
 
 /**
- * How long "no entitlement" must hold before a running Enterprise deployment believes it: longer
- * than an unreturned unit lasts, plus one re-check, so a unit held elsewhere can never cause it.
+ * How long "no entitlement" must hold, with no other answer in between, before a running Enterprise
+ * deployment believes it: longer than an unreturned unit lasts, plus one re-check, so a unit held
+ * elsewhere can never cause it.
  */
 const NO_ENTITLEMENT_CONFIRMATION_MS = UNRETURNED_UNIT_MS + RECHECK_INTERVAL_MS
 
@@ -211,10 +212,10 @@ export async function applyAwsMarketplaceContainerLicense(): Promise<void> {
   }
   const recheck = async () => {
     const state = await checkLicense(client, productSku)
-    if (state === 'entitled') {
-      noEntitlementSince = null
-      if (!enterpriseAtBoot) return restartToApply(state)
-    } else if (state === 'no-entitlement' && enterpriseAtBoot) {
+    // Only an unbroken run of "no entitlement" answers counts toward switching Enterprise off.
+    if (state !== 'no-entitlement') noEntitlementSince = null
+    if (state === 'entitled' && !enterpriseAtBoot) return restartToApply(state)
+    if (state === 'no-entitlement' && enterpriseAtBoot) {
       noEntitlementSince ??= Date.now()
       if (Date.now() - noEntitlementSince >= NO_ENTITLEMENT_CONFIRMATION_MS) {
         return restartToApply(state)

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