ai-functions/test/context-isolation.test.ts

Core AI primitives for building intelligent applications

/**
 * Tests for Concurrent Context Isolation
 *
 * TDD: RED phase - These tests expose race conditions in global context management
 *
 * The current implementation uses a global `defaultClient` variable that can cause
 * context bleeding between concurrent operations. These tests demonstrate:
 *
 * 1. Context leakage between concurrent Promise.all operations
 * 2. Configuration changes affecting in-flight requests
 * 3. API key isolation failure in multi-tenant scenarios
 * 4. Async/await interleaving causing wrong context
 *
 * Expected: These tests should FAIL or be FLAKY with the current implementation,
 * demonstrating the need for AsyncLocalStorage-based context isolation.
 */

import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'
import {
  configure,
  resetContext,
  withContext,
  getContext,
  getModel,
  getProvider,
  getBatchMode,
  type ExecutionContext,
} from '../src/context.js'

// ============================================================================
// Test Utilities
// ============================================================================

/**
 * Simulates an async AI operation that takes variable time
 * This helps expose race conditions by introducing realistic delays
 */
async function simulateAICall(delayMs: number): Promise<ExecutionContext> {
  // Capture context at start
  const startContext = getContext()

  // Simulate network delay
  await new Promise(resolve => setTimeout(resolve, delayMs))

  // Capture context at end - should be same as start!
  const endContext = getContext()

  return {
    // Return both for comparison
    ...endContext,
    _startModel: startContext.model,
    _startProvider: startContext.provider,
  } as ExecutionContext & { _startModel?: string; _startProvider?: string }
}

/**
 * Creates a delayed operation that should maintain its context
 */
function createDelayedContextCapture(delayMs: number) {
  return async (): Promise<{ model: string | undefined; provider: string | undefined }> => {
    const beforeDelay = { model: getModel(), provider: getProvider() }
    await new Promise(resolve => setTimeout(resolve, delayMs))
    const afterDelay = { model: getModel(), provider: getProvider() }

    // These should be equal - if not, context leaked!
    return {
      model: afterDelay.model,
      provider: afterDelay.provider,
      _beforeModel: beforeDelay.model,
      _beforeProvider: beforeDelay.provider,
    } as { model: string | undefined; provider: string | undefined }
  }
}

// ============================================================================
// Context Isolation Tests
// ============================================================================

describe('Concurrent Context Isolation', () => {
  beforeEach(() => {
    resetContext()
  })

  afterEach(() => {
    resetContext()
  })

  describe('Promise.all Context Leakage', () => {
    it('should isolate context between concurrent withContext calls', async () => {
      // This test exposes the race condition in the fallback implementation
      // where global context is temporarily modified and restored
      const results = await Promise.all([
        withContext({ model: 'claude-opus-4-5', provider: 'anthropic' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 50))
          return { model: getModel(), provider: getProvider() }
        }),
        withContext({ model: 'gpt-4o', provider: 'openai' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 10))
          return { model: getModel(), provider: getProvider() }
        }),
        withContext({ model: 'gemini-pro', provider: 'google' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 30))
          return { model: getModel(), provider: getProvider() }
        }),
      ])

      // Each operation should see its own context, not another's
      expect(results[0]).toEqual({ model: 'claude-opus-4-5', provider: 'anthropic' })
      expect(results[1]).toEqual({ model: 'gpt-4o', provider: 'openai' })
      expect(results[2]).toEqual({ model: 'gemini-pro', provider: 'google' })
    })

    it('should not leak context when operations complete in different order', async () => {
      // Operations with different completion times
      const contexts: Array<{ model: string | undefined; provider: string | undefined; order: number }> = []

      await Promise.all([
        withContext({ model: 'slow-model', provider: 'anthropic' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 100))
          contexts.push({ model: getModel(), provider: getProvider(), order: 1 })
        }),
        withContext({ model: 'fast-model', provider: 'openai' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 10))
          contexts.push({ model: getModel(), provider: getProvider(), order: 2 })
        }),
      ])

      // Fast model completes first (order 2), slow model second (order 1)
      // But each should see its own model
      const slowResult = contexts.find(c => c.order === 1)
      const fastResult = contexts.find(c => c.order === 2)

      expect(slowResult?.model).toBe('slow-model')
      expect(fastResult?.model).toBe('fast-model')
    })

    it('should handle nested withContext calls concurrently', async () => {
      const results = await Promise.all([
        withContext({ model: 'outer-1' }, async () => {
          const outerModel = getModel()
          const innerResult = await withContext({ model: 'inner-1' }, async () => {
            await new Promise(resolve => setTimeout(resolve, 20))
            return getModel()
          })
          // After inner completes, should restore outer context
          const afterInnerModel = getModel()
          return { outerModel, innerResult, afterInnerModel }
        }),
        withContext({ model: 'outer-2' }, async () => {
          const outerModel = getModel()
          await new Promise(resolve => setTimeout(resolve, 10))
          return { outerModel, innerResult: null, afterInnerModel: getModel() }
        }),
      ])

      expect(results[0].outerModel).toBe('outer-1')
      expect(results[0].innerResult).toBe('inner-1')
      expect(results[0].afterInnerModel).toBe('outer-1')
      expect(results[1].outerModel).toBe('outer-2')
      expect(results[1].afterInnerModel).toBe('outer-2')
    })
  })

  describe('API Key Isolation (Multi-tenant)', () => {
    it('should isolate API keys between concurrent tenant requests', async () => {
      // Simulate multi-tenant scenario where each request has different credentials
      interface TenantContext {
        apiKey: string
        tenantId: string
      }

      const tenantResults: Array<TenantContext & { seenApiKey: string }> = []

      await Promise.all([
        withContext({ metadata: { apiKey: 'tenant-1-key', tenantId: 'tenant-1' } }, async () => {
          await new Promise(resolve => setTimeout(resolve, 50))
          const ctx = getContext()
          tenantResults.push({
            apiKey: 'tenant-1-key',
            tenantId: 'tenant-1',
            seenApiKey: (ctx.metadata as TenantContext)?.apiKey,
          })
        }),
        withContext({ metadata: { apiKey: 'tenant-2-key', tenantId: 'tenant-2' } }, async () => {
          await new Promise(resolve => setTimeout(resolve, 10))
          const ctx = getContext()
          tenantResults.push({
            apiKey: 'tenant-2-key',
            tenantId: 'tenant-2',
            seenApiKey: (ctx.metadata as TenantContext)?.apiKey,
          })
        }),
        withContext({ metadata: { apiKey: 'tenant-3-key', tenantId: 'tenant-3' } }, async () => {
          await new Promise(resolve => setTimeout(resolve, 30))
          const ctx = getContext()
          tenantResults.push({
            apiKey: 'tenant-3-key',
            tenantId: 'tenant-3',
            seenApiKey: (ctx.metadata as TenantContext)?.apiKey,
          })
        }),
      ])

      // Each tenant should see their own API key, not another tenant's!
      // This is a critical security issue if context leaks
      for (const result of tenantResults) {
        expect(result.seenApiKey).toBe(result.apiKey)
      }
    })
  })

  describe('Configuration Changes Mid-flight', () => {
    it('should not affect in-flight requests when global config changes', async () => {
      configure({ model: 'initial-model', provider: 'anthropic' })

      // Start a long-running operation
      const operationPromise = withContext({}, async () => {
        const startModel = getModel()
        await new Promise(resolve => setTimeout(resolve, 100))
        const endModel = getModel()
        return { startModel, endModel }
      })

      // Change global config while operation is in flight
      await new Promise(resolve => setTimeout(resolve, 10))
      configure({ model: 'changed-model', provider: 'openai' })

      const result = await operationPromise

      // The operation should see consistent context throughout
      // Even though global config changed mid-flight
      expect(result.startModel).toBe(result.endModel)
      expect(result.startModel).toBe('initial-model')
    })

    it.skip('should isolate configure() calls from concurrent operations (known limitation - use withContext)', async () => {
      const results: string[] = []

      await Promise.all([
        // Operation 1: Set config and use it
        (async () => {
          configure({ model: 'op1-model' })
          await new Promise(resolve => setTimeout(resolve, 50))
          results.push(`op1: ${getModel()}`)
        })(),
        // Operation 2: Set different config
        (async () => {
          await new Promise(resolve => setTimeout(resolve, 10))
          configure({ model: 'op2-model' })
          await new Promise(resolve => setTimeout(resolve, 10))
          results.push(`op2: ${getModel()}`)
        })(),
      ])

      // With proper isolation, each operation would see its own config
      // With current implementation, they interfere with each other
      // The test documents the expected behavior vs actual behavior
      expect(results).toContain('op1: op1-model')
      expect(results).toContain('op2: op2-model')
    })
  })

  describe('Async/Await Interleaving', () => {
    it('should maintain context across await points', async () => {
      const results: Array<{ step: string; model: string | undefined }> = []

      await withContext({ model: 'test-model' }, async () => {
        results.push({ step: 'before-await-1', model: getModel() })
        await new Promise(resolve => setTimeout(resolve, 10))

        results.push({ step: 'after-await-1', model: getModel() })
        await new Promise(resolve => setTimeout(resolve, 10))

        results.push({ step: 'after-await-2', model: getModel() })
        await new Promise(resolve => setTimeout(resolve, 10))

        results.push({ step: 'after-await-3', model: getModel() })
      })

      // All steps should see the same model
      for (const result of results) {
        expect(result.model).toBe('test-model')
      }
    })

    it('should handle interleaved async generators', async () => {
      async function* contextAwareGenerator(contextModel: string, steps: number) {
        for (let i = 0; i < steps; i++) {
          await new Promise(resolve => setTimeout(resolve, Math.random() * 20))
          yield { step: i, model: getModel(), expected: contextModel }
        }
      }

      const results: Array<{ step: number; model: string | undefined; expected: string }> = []

      // Run generators concurrently in different contexts
      await Promise.all([
        withContext({ model: 'generator-1' }, async () => {
          for await (const item of contextAwareGenerator('generator-1', 5)) {
            results.push(item)
          }
        }),
        withContext({ model: 'generator-2' }, async () => {
          for await (const item of contextAwareGenerator('generator-2', 5)) {
            results.push(item)
          }
        }),
      ])

      // Each generator should see its own context
      for (const result of results) {
        expect(result.model).toBe(result.expected)
      }
    })
  })

  describe('Batch Mode Isolation', () => {
    it('should isolate batch mode settings between concurrent operations', async () => {
      const results = await Promise.all([
        withContext({ batchMode: 'immediate' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 30))
          return getBatchMode()
        }),
        withContext({ batchMode: 'deferred' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 10))
          return getBatchMode()
        }),
        withContext({ batchMode: 'flex' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 20))
          return getBatchMode()
        }),
      ])

      expect(results[0]).toBe('immediate')
      expect(results[1]).toBe('deferred')
      expect(results[2]).toBe('flex')
    })
  })

  describe('High Concurrency Stress Test', () => {
    it('should maintain context isolation under high concurrency', async () => {
      const concurrency = 50
      const operations = Array.from({ length: concurrency }, (_, i) => i)

      const results = await Promise.all(
        operations.map(i =>
          withContext({ model: `model-${i}`, metadata: { opId: i } }, async () => {
            // Random delay to maximize interleaving
            await new Promise(resolve => setTimeout(resolve, Math.random() * 50))
            const ctx = getContext()
            return {
              expected: `model-${i}`,
              actual: ctx.model,
              opId: (ctx.metadata as { opId: number })?.opId,
            }
          })
        )
      )

      // Every operation should see its own context
      let failures = 0
      for (const result of results) {
        if (result.expected !== result.actual) {
          failures++
        }
      }

      // With proper isolation, there should be 0 failures
      // With the current global-based fallback, we expect failures
      expect(failures).toBe(0)
    })

    it('should handle rapid context switches without corruption', async () => {
      const iterations = 100
      const errors: string[] = []

      for (let i = 0; i < iterations; i++) {
        const expectedModel = `rapid-${i}`

        await withContext({ model: expectedModel }, async () => {
          const actualModel = getModel()
          if (actualModel !== expectedModel) {
            errors.push(`Iteration ${i}: expected ${expectedModel}, got ${actualModel}`)
          }
        })
      }

      expect(errors).toEqual([])
    })
  })

  describe('Context Restoration After Error', () => {
    it('should restore context after exception in withContext', async () => {
      configure({ model: 'original-model' })

      try {
        await withContext({ model: 'error-model' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 10))
          throw new Error('Test error')
        })
      } catch {
        // Expected error
      }

      // Context should be restored to original
      expect(getModel()).toBe('original-model')
    })

    it('should restore context even when nested contexts throw', async () => {
      configure({ model: 'root-model' })

      try {
        await withContext({ model: 'outer-model' }, async () => {
          await withContext({ model: 'inner-model' }, async () => {
            await new Promise(resolve => setTimeout(resolve, 10))
            throw new Error('Inner error')
          })
        })
      } catch {
        // Expected error
      }

      // Context should be fully restored
      expect(getModel()).toBe('root-model')
    })

    it('should isolate error context from parallel operations', async () => {
      const results: Array<{ id: number; model: string | undefined; error?: boolean }> = []

      await Promise.allSettled([
        withContext({ model: 'success-1' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 50))
          results.push({ id: 1, model: getModel() })
        }),
        withContext({ model: 'error-op' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 20))
          throw new Error('Deliberate error')
        }),
        withContext({ model: 'success-2' }, async () => {
          await new Promise(resolve => setTimeout(resolve, 30))
          results.push({ id: 2, model: getModel() })
        }),
      ])

      // The error in one operation should not affect others
      const op1 = results.find(r => r.id === 1)
      const op2 = results.find(r => r.id === 2)

      expect(op1?.model).toBe('success-1')
      expect(op2?.model).toBe('success-2')
    })
  })
})

// ============================================================================
// Global Configure Race Conditions
// ============================================================================

describe('Global Configure Race Conditions', () => {
  beforeEach(() => {
    resetContext()
  })

  afterEach(() => {
    resetContext()
  })

  // NOTE: These tests document the EXPECTED race conditions when using configure()
  // concurrently. The configure() function is for setting global defaults at startup,
  // NOT for concurrent tenant isolation. Use withContext() for that.
  //
  // These tests are marked as .skip because they document known limitations,
  // not bugs. The correct fix is to use withContext() for concurrent scenarios.

  it.skip('should not allow configure() to affect already-started operations (known limitation - use withContext)', async () => {
    // This is the core race condition: configure() modifies globalContext
    // which is read by getContext() at arbitrary times

    configure({ model: 'original' })

    const operationResults: Array<{ stage: string; model: string | undefined }> = []

    // Start operation that reads context multiple times
    const operation = (async () => {
      operationResults.push({ stage: 'start', model: getModel() })
      await new Promise(resolve => setTimeout(resolve, 30))
      operationResults.push({ stage: 'middle', model: getModel() })
      await new Promise(resolve => setTimeout(resolve, 30))
      operationResults.push({ stage: 'end', model: getModel() })
    })()

    // Interfering configure() calls while operation is running
    await new Promise(resolve => setTimeout(resolve, 10))
    configure({ model: 'interference-1' })
    await new Promise(resolve => setTimeout(resolve, 20))
    configure({ model: 'interference-2' })
    await new Promise(resolve => setTimeout(resolve, 20))
    configure({ model: 'interference-3' })

    await operation

    // The operation should see consistent context throughout
    // But with global configure(), each stage might see different models!
    const models = operationResults.map(r => r.model)
    const uniqueModels = new Set(models)

    // If isolation is working, should only see 1 model
    // If broken, might see multiple models
    expect(uniqueModels.size).toBe(1)
    expect(models[0]).toBe('original')
  })

  it.skip('should demonstrate configure() race in multi-tenant scenario (known limitation - use withContext)', async () => {
    // Simulate multiple tenants configuring their own settings
    // This is a realistic scenario that will fail without proper isolation

    const tenantOperations = [
      { tenantId: 'tenant-a', model: 'claude-opus-4-5', delay: 100 },
      { tenantId: 'tenant-b', model: 'gpt-4o', delay: 50 },
      { tenantId: 'tenant-c', model: 'gemini-pro', delay: 75 },
    ]

    const results: Array<{ tenantId: string; expected: string; actual: string | undefined }> = []

    await Promise.all(
      tenantOperations.map(async ({ tenantId, model, delay }) => {
        // Each tenant configures their model
        configure({ model })

        // Some processing time...
        await new Promise(resolve => setTimeout(resolve, delay))

        // Check what model they see - should be their configured model
        const actualModel = getModel()
        results.push({ tenantId, expected: model, actual: actualModel })
      })
    )

    // Each tenant should see their own model
    // This WILL fail because configure() races with other tenants
    for (const result of results) {
      expect(result.actual).toBe(result.expected)
    }
  })

  it('should show configure() racing with withContext()', async () => {
    configure({ model: 'global-model' })

    const results: string[] = []

    await Promise.all([
      // Operation using withContext (should be isolated)
      withContext({ model: 'isolated-model' }, async () => {
        results.push(`withContext-start: ${getModel()}`)
        await new Promise(resolve => setTimeout(resolve, 50))
        results.push(`withContext-end: ${getModel()}`)
      }),
      // Operation modifying global config (affects anyone reading globalContext)
      (async () => {
        await new Promise(resolve => setTimeout(resolve, 25))
        configure({ model: 'racing-model' })
        results.push(`configure-done: ${getModel()}`)
      })(),
    ])

    // withContext should maintain isolation
    expect(results).toContain('withContext-start: isolated-model')
    expect(results).toContain('withContext-end: isolated-model')
  })

  it.skip('should expose getGlobalContext() mutation issues (known limitation - use withContext)', async () => {
    configure({ model: 'safe-model' })

    // Multiple operations reading global context
    const reads: Array<{ time: number; model: string | undefined }> = []

    const readerPromise = (async () => {
      for (let i = 0; i < 10; i++) {
        const ctx = getContext()
        reads.push({ time: Date.now(), model: ctx.model })
        await new Promise(resolve => setTimeout(resolve, 10))
      }
    })()

    // Concurrent writer
    const writerPromise = (async () => {
      for (let i = 0; i < 5; i++) {
        await new Promise(resolve => setTimeout(resolve, 15))
        configure({ model: `mutated-${i}` })
      }
    })()

    await Promise.all([readerPromise, writerPromise])

    // The reader should see consistent values if properly isolated
    // With current implementation, it will see a mix of values
    const uniqueModels = new Set(reads.map(r => r.model))

    // Document the race condition - multiple values seen
    // With proper isolation, should only see initial value
    console.log('Models seen during concurrent read/write:', [...uniqueModels])

    // This assertion documents expected behavior - only 1 model should be seen
    expect(uniqueModels.size).toBe(1)
  })
})

// ============================================================================
// Edge Cases and Regression Tests
// ============================================================================

describe('Context Edge Cases', () => {
  beforeEach(() => {
    resetContext()
  })

  afterEach(() => {
    resetContext()
  })

  it('should handle undefined context values', async () => {
    const results = await Promise.all([
      withContext({ model: undefined }, async () => {
        await new Promise(resolve => setTimeout(resolve, 10))
        return getModel()
      }),
      withContext({ model: 'defined-model' }, async () => {
        await new Promise(resolve => setTimeout(resolve, 5))
        return getModel()
      }),
    ])

    // Default model should be returned for undefined
    expect(results[0]).toBe('sonnet')
    expect(results[1]).toBe('defined-model')
  })

  it('should handle empty context object', async () => {
    configure({ model: 'configured-model' })

    const result = await withContext({}, async () => {
      await new Promise(resolve => setTimeout(resolve, 10))
      return getModel()
    })

    // Should inherit from global config
    expect(result).toBe('configured-model')
  })

  it('should handle deeply nested concurrent contexts', async () => {
    const results: string[] = []

    await Promise.all([
      withContext({ model: 'level-0-a' }, async () => {
        results.push(`0a: ${getModel()}`)
        await withContext({ model: 'level-1-a' }, async () => {
          results.push(`1a: ${getModel()}`)
          await Promise.all([
            withContext({ model: 'level-2-a1' }, async () => {
              await new Promise(resolve => setTimeout(resolve, 20))
              results.push(`2a1: ${getModel()}`)
            }),
            withContext({ model: 'level-2-a2' }, async () => {
              await new Promise(resolve => setTimeout(resolve, 10))
              results.push(`2a2: ${getModel()}`)
            }),
          ])
          results.push(`1a-after: ${getModel()}`)
        })
        results.push(`0a-after: ${getModel()}`)
      }),
      withContext({ model: 'level-0-b' }, async () => {
        await new Promise(resolve => setTimeout(resolve, 15))
        results.push(`0b: ${getModel()}`)
      }),
    ])

    // Verify each level sees correct context
    expect(results).toContain('0a: level-0-a')
    expect(results).toContain('1a: level-1-a')
    expect(results).toContain('2a1: level-2-a1')
    expect(results).toContain('2a2: level-2-a2')
    expect(results).toContain('1a-after: level-1-a')
    expect(results).toContain('0a-after: level-0-a')
    expect(results).toContain('0b: level-0-b')
  })
})