@@ -111,7 +111,9 @@ const {
111111 mockExecuteChatFunction,
112112 mockGetInternalToolOperationHandler,
113113 mockExecuteInternalToolOperation,
114+ mockPayerHasPaidAccess,
114115} = vi . hoisted ( ( ) => ( {
116+ mockPayerHasPaidAccess : vi . fn ( ) ,
115117 mockGetToolAsync : vi . fn ( ) ,
116118 mockRateLimiterFns : {
117119 acquireKey : vi . fn ( ) ,
@@ -173,6 +175,10 @@ vi.mock('@/lib/core/rate-limiter/hosted-key', () => ({
173175 getHostedKeyRateLimiter : ( ) => mockRateLimiterFns ,
174176} ) )
175177
178+ vi . mock ( '@/lib/billing/core/chat-access' , ( ) => ( {
179+ payerHasPaidAccess : mockPayerHasPaidAccess ,
180+ } ) )
181+
176182vi . mock ( '@/lib/uploads/contexts/workspace/file-manager' , ( ) => workspaceFileManagerMock )
177183
178184vi . mock (
@@ -6551,6 +6557,97 @@ describe('Hosted Key Injection', () => {
65516557 }
65526558 } )
65536559
6560+ describe ( 'requiresPaidPlan' , ( ) => {
6561+ const PAID_ONLY_TOOL_ID = 'test_paid_only_hosting'
6562+
6563+ beforeEach ( ( ) => {
6564+ setEnvFlags ( { isHosted : true , isBillingEnabled : true } )
6565+ mockRateLimiterFns . acquireKey . mockResolvedValue ( {
6566+ success : true ,
6567+ key : 'platform-fixture-key' ,
6568+ keyIndex : 0 ,
6569+ envVarName : 'TEST_HOSTED_KEY' ,
6570+ } )
6571+ tools [ PAID_ONLY_TOOL_ID ] = {
6572+ id : PAID_ONLY_TOOL_ID ,
6573+ name : 'Paid-only fixture' ,
6574+ description : 'Hosted key reserved for paying payers' ,
6575+ version : '1.0.0' ,
6576+ params : { } ,
6577+ request : { url : 'https://example.com/fixture' , method : 'GET' , headers : ( ) => ( { } ) } ,
6578+ hosting : {
6579+ envKeyPrefix : 'TEST_HOSTED_KEY' ,
6580+ apiKeyParam : 'apiKey' ,
6581+ byokProviderId : 'falai' ,
6582+ requiresPaidPlan : true ,
6583+ pricing : { type : 'per_request' , cost : 0.5 } ,
6584+ rateLimit : { mode : 'per_request' , requestsPerMinute : 100 } ,
6585+ } ,
6586+ transformResponse : async ( ) => ( { success : true , output : { } } ) ,
6587+ }
6588+ } )
6589+
6590+ afterEach ( ( ) => {
6591+ Reflect . deleteProperty ( tools , PAID_ONLY_TOOL_ID )
6592+ resetEnvFlagsMock ( )
6593+ } )
6594+
6595+ async function runPaidOnlyTool ( ) {
6596+ mockValidateUrlWithDNS . mockResolvedValue ( { isValid : true , resolvedIP : '93.184.216.34' } )
6597+ mockSecureFetchWithPinnedIP . mockResolvedValue ( {
6598+ ok : true ,
6599+ status : 200 ,
6600+ statusText : 'OK' ,
6601+ headers : { get : ( ) => null , toRecord : ( ) => ( { } ) } ,
6602+ json : async ( ) => ( { ok : true } ) ,
6603+ text : async ( ) => '{"ok":true}' ,
6604+ } )
6605+ const request = mockSecureFetchWithPinnedIP
6606+ const result = await executeTool (
6607+ PAID_ONLY_TOOL_ID ,
6608+ { } ,
6609+ { executionContext : createToolExecutionContext ( { userId : 'user-1' } ) }
6610+ )
6611+ return { result, request }
6612+ }
6613+
6614+ it ( 'refuses the hosted key to a payer without paid access before acquiring one' , async ( ) => {
6615+ mockGetBYOKKey . mockResolvedValue ( null )
6616+ mockPayerHasPaidAccess . mockResolvedValue ( false )
6617+
6618+ const { result, request } = await runPaidOnlyTool ( )
6619+
6620+ expect ( result . success ) . toBe ( false )
6621+ expect ( result . error ) . toContain ( 'available on paid plans' )
6622+ expect ( mockPayerHasPaidAccess ) . toHaveBeenCalledWith ( TEST_BILLING_ATTRIBUTION )
6623+ expect ( mockRateLimiterFns . acquireKey ) . not . toHaveBeenCalled ( )
6624+ expect ( request ) . not . toHaveBeenCalled ( )
6625+ } )
6626+
6627+ it ( 'grants the hosted key to a payer with paid access' , async ( ) => {
6628+ mockGetBYOKKey . mockResolvedValue ( null )
6629+ mockPayerHasPaidAccess . mockResolvedValue ( true )
6630+
6631+ const { result, request } = await runPaidOnlyTool ( )
6632+
6633+ expect ( result . success ) . toBe ( true )
6634+ expect ( mockRateLimiterFns . acquireKey ) . toHaveBeenCalledTimes ( 1 )
6635+ expect ( request ) . toHaveBeenCalledTimes ( 1 )
6636+ } )
6637+
6638+ it ( "runs a free payer's own workspace key without consulting the plan" , async ( ) => {
6639+ mockGetBYOKKey . mockResolvedValue ( { apiKey : 'workspace-own-key' } )
6640+ mockPayerHasPaidAccess . mockResolvedValue ( false )
6641+
6642+ const { result, request } = await runPaidOnlyTool ( )
6643+
6644+ expect ( result . success ) . toBe ( true )
6645+ expect ( mockPayerHasPaidAccess ) . not . toHaveBeenCalled ( )
6646+ expect ( mockRateLimiterFns . acquireKey ) . not . toHaveBeenCalled ( )
6647+ expect ( request ) . toHaveBeenCalledTimes ( 1 )
6648+ } )
6649+ } )
6650+
65546651 it ( 'should not inject hosted key when tool has no hosting config' , async ( ) => {
65556652 const mockTool = {
65566653 id : 'test_no_hosting' ,
0 commit comments