feat: cache AccessDenied error for count tokens (#1032)

strands-ts/src/models/__tests__/bedrock.test.ts

@@ -4294,6 +4294,26 @@ describe('BedrockModel', () => {
       expect(mockSend).toHaveBeenCalledOnce()
     })
 
+    it('should cache model ID and skip API call on AccessDeniedException', async () => {
+      const accessDeniedError = new Error(
+        'User: arn:aws:sts::123456789012:assumed-role/role is not authorized to perform: bedrock:CountTokens'
+      )
+      accessDeniedError.name = 'AccessDeniedException'
+      const mockSend = vi.fn(async () => {
+        throw accessDeniedError
+      })
+      mockBedrockClientImplementation({ send: mockSend })
+      const model = new BedrockModel()
+
+      // First call: hits API, gets AccessDeniedException, caches
+      await model.countTokens(messages)
+      expect(mockSend).toHaveBeenCalledOnce()
+
+      // Second call: skips API entirely due to caching
+      await model.countTokens(messages)
+      expect(mockSend).toHaveBeenCalledOnce()
+    })
+
     it('should not cache model ID for other errors', async () => {
       const mockSend = vi.fn(async () => {
         throw new Error('Transient network error')

strands-ts/src/models/bedrock.ts

@@ -97,10 +97,10 @@ const BEDROCK_CONTEXT_WINDOW_OVERFLOW_MESSAGES = [
 ]
 
 /**
- * Cache of model IDs that do not support the CountTokens API.
- * Prevents repeated failing API calls for models that will never support token counting.
+ * Cache of model IDs for which CountTokens API calls should be skipped.
+ * Prevents repeated failing API calls that will never succeed for the lifetime of the process.
  */
-const UNSUPPORTED_COUNT_TOKENS_MODELS = new Set<string>()
+const SKIP_COUNT_TOKENS_MODELS = new Set<string>()
 
 /**
  * Mapping of Bedrock stop reasons to SDK stop reasons.
@@ -349,13 +349,13 @@ export class BedrockModel extends Model<BedrockModelConfig> {
   private _client: BedrockRuntimeClient
 
   /**
-   * Clears the cache of model IDs that do not support the CountTokens API.
+   * Clears the cache of model IDs for which CountTokens is skipped.
    * After calling this, the next countTokens invocation will attempt the API again.
    *
    * @internal
    */
   static clearCountTokensCache(): void {
-    UNSUPPORTED_COUNT_TOKENS_MODELS.clear()
+    SKIP_COUNT_TOKENS_MODELS.clear()
   }
 
   /**
@@ -514,7 +514,7 @@ export class BedrockModel extends Model<BedrockModelConfig> {
 
     const modelId = this._config.modelId ?? MODEL_DEFAULTS.bedrock.modelId
 
-    if (UNSUPPORTED_COUNT_TOKENS_MODELS.has(modelId)) {
+    if (SKIP_COUNT_TOKENS_MODELS.has(modelId)) {
       return super.countTokens(messages, options)
     }
 
@@ -539,7 +539,13 @@ export class BedrockModel extends Model<BedrockModelConfig> {
       logger.debug(`total_tokens=<${response.inputTokens}> | native token count`)
       return response.inputTokens
     } catch (error) {
-      if (
+      if (error instanceof Error && error.name === 'AccessDeniedException') {
+        warnOnce(
+          logger,
+          `model_id=<${modelId}> | bedrock:CountTokens permission denied, falling back to heuristic estimation`
+        )
+        SKIP_COUNT_TOKENS_MODELS.add(modelId)
+      } else if (
         error instanceof Error &&
         error.name === 'ValidationException' &&
         error.message.includes("doesn't support counting tokens")
@@ -547,7 +553,7 @@ export class BedrockModel extends Model<BedrockModelConfig> {
         logger.debug(
           `model_id=<${modelId}> | model does not support CountTokens, caching for future calls, falling back to estimation`
         )
-        UNSUPPORTED_COUNT_TOKENS_MODELS.add(modelId)
+        SKIP_COUNT_TOKENS_MODELS.add(modelId)
       } else {
         logger.debug(`error=<${error}> | native token counting failed, falling back to estimation`)
       }