diff --git a/packages/tenants/src/context/TenantAgentContextProvider.ts b/packages/tenants/src/context/TenantAgentContextProvider.ts index 5205475339..3925179e38 100644 --- a/packages/tenants/src/context/TenantAgentContextProvider.ts +++ b/packages/tenants/src/context/TenantAgentContextProvider.ts @@ -179,11 +179,14 @@ export class TenantAgentContextProvider implements AgentContextProvider { * to update the storage for a tenant manually */ public async updateTenantStorage(tenantRecord: TenantRecord, updateOptions?: UpdateAssistantUpdateOptions) { - await this.tenantSessionCoordinator.getContextForSession(tenantRecord, { + const agentContext = await this.tenantSessionCoordinator.getContextForSession(tenantRecord, { // runInMutex allows us to run the updateTenantStorage method in a mutex lock // prevent other sessions from being started while the update is in progress runInMutex: (agentContext) => this._updateTenantStorage(tenantRecord, agentContext, updateOptions), }) + + // End sesion afterwards + await agentContext.endSession() } /** diff --git a/packages/tenants/src/context/TenantSessionCoordinator.ts b/packages/tenants/src/context/TenantSessionCoordinator.ts index b7d81ff70b..466ebd1a66 100644 --- a/packages/tenants/src/context/TenantSessionCoordinator.ts +++ b/packages/tenants/src/context/TenantSessionCoordinator.ts @@ -54,6 +54,10 @@ export class TenantSessionCoordinator { ) } + public getSessionCountForTenant(tenantId: string) { + return this.tenantAgentContextMapping[tenantId]?.sessionCount ?? 0 + } + /** * Get agent context to use for a session. If an agent context for this tenant does not exist yet * it will create it and store it for later use. If the agent context does already exist it will @@ -94,7 +98,23 @@ export class TenantSessionCoordinator { ) if (runInMutex) { - await runInMutex(tenantSessions.agentContext) + try { + await runInMutex(tenantSessions.agentContext) + } catch (error) { + // If the runInMutex failed we should release the session again + tenantSessions.sessionCount-- + this.logger.debug( + `Decreased agent context session count for tenant '${tenantSessions.agentContext.contextCorrelationId}' to ${tenantSessions.sessionCount} due to failure in mutex script`, + error + ) + + if (tenantSessions.sessionCount <= 0 && tenantSessions.agentContext) { + await this.closeAgentContext(tenantSessions.agentContext) + delete this.tenantAgentContextMapping[tenantSessions.agentContext.contextCorrelationId] + } + + throw error + } } return tenantSessions.agentContext diff --git a/packages/tenants/tests/tenants-storage-update.test.ts b/packages/tenants/tests/tenants-storage-update.test.ts index e26b6ae421..28eb14293f 100644 --- a/packages/tenants/tests/tenants-storage-update.test.ts +++ b/packages/tenants/tests/tenants-storage-update.test.ts @@ -14,6 +14,7 @@ import path from 'path' import { AskarModule, AskarMultiWalletDatabaseScheme } from '../../askar/src' import { ariesAskar } from '../../askar/tests/helpers' import { testLogger } from '../../core/tests' +import { TenantSessionCoordinator } from '../src/context/TenantSessionCoordinator' import { TenantsModule } from '@credo-ts/tenants' @@ -193,6 +194,9 @@ describe('Tenants Storage Update', () => { agent.modules.tenants.getTenantAgent({ tenantId: '1d45d3c2-3480-4375-ac6f-47c322f091b0' }) ).rejects.toThrow(/Current agent storage for tenant 1d45d3c2-3480-4375-ac6f-47c322f091b0 is not up to date/) + const tenantSessionCoordinator = agent.dependencyManager.resolve(TenantSessionCoordinator) + expect(tenantSessionCoordinator.getSessionCountForTenant(tenant.id)).toBe(0) + // Update tenant await agent.modules.tenants.updateTenantStorage({ tenantId: tenant.id, @@ -201,6 +205,9 @@ describe('Tenants Storage Update', () => { }, }) + // Should have closed session after upgrade + expect(tenantSessionCoordinator.getSessionCountForTenant(tenant.id)).toBe(0) + // Expect tenant storage version to be 0.5 const updatedTenant = await agent.modules.tenants.getTenantById('1d45d3c2-3480-4375-ac6f-47c322f091b0') expect(updatedTenant.storageVersion).toBe('0.5')