@lenne.tech/nest-server 12.0.0 → 12.1.0

package/src/core/modules/ai/core-ai.constants.ts

@@ -90,3 +90,13 @@ export const AI_TOOL_GRANT_MODEL = 'AiToolGrant';
 export const AI_TOOL_POLICY_CLASS = 'AI_TOOL_POLICY_CLASS';
 /** Mongoose injection token. */
 export const AI_TOOL_POLICY_MODEL = 'AiToolPolicy';
+
+// ---- MCP OAuth (stored natively on MongoDB; these models give them their tables on a SQL source) ----
+/** Mongoose model name of the registered OAuth clients (`aiMcpOAuthClients`). */
+export const AI_MCP_OAUTH_CLIENT_MODEL = 'AiMcpOAuthClient';
+/** Mongoose model name of the authorization codes (`aiMcpOAuthCodes`). */
+export const AI_MCP_OAUTH_CODE_MODEL = 'AiMcpOAuthCode';
+/** Mongoose model name of the refresh tokens (`aiMcpOAuthRefreshTokens`). */
+export const AI_MCP_OAUTH_REFRESH_TOKEN_MODEL = 'AiMcpOAuthRefreshToken';
+/** How long an authorization code may be exchanged — checked on every read, not left to a TTL index. */
+export const AI_MCP_OAUTH_CODE_TTL_SECONDS = 600;

package/src/core/modules/ai/core-ai.module.ts

@@ -1,5 +1,4 @@
 import { DynamicModule, Global, Module, Type } from '@nestjs/common';
-import { MongooseModule } from '@nestjs/mongoose';
 
 import { CoreAiMcpOAuthInitializer } from './core-ai-mcp-oauth.initializer.js';
 import { CoreAiMcpController } from './core-ai-mcp.controller.js';
@@ -7,7 +6,18 @@ import { CoreAiController } from './core-ai.controller.js';
 import { CoreAiResolver } from './core-ai.resolver.js';
 import { CoreAiMcpOAuthService } from './services/core-ai-mcp-oauth.service.js';
 import { CoreAiMcpService } from './services/core-ai-mcp.service.js';
+import { CorePersistenceModule } from '../../common/persistence/core-persistence.module.js';
+import {
+  AI_MCP_OAUTH_CLIENT_MODEL,
+  AI_MCP_OAUTH_CODE_MODEL,
+  AI_MCP_OAUTH_REFRESH_TOKEN_MODEL,
+} from './core-ai.constants.js';
 import { AiBudgetLimitSchema, CoreAiBudgetLimit } from './models/core-ai-budget-limit.model.js';
+import {
+  AiMcpOAuthClientSchema,
+  AiMcpOAuthCodeSchema,
+  AiMcpOAuthRefreshTokenSchema,
+} from './models/core-ai-mcp-oauth.models.js';
 import {
   AiConnectionPreferenceSchema,
   CoreAiConnectionPreference,
@@ -184,6 +194,25 @@ export class CoreAiModule {
     const SlotServiceClass = options.slotService || CoreAiSlotService;
     const ResolverClass = options.resolver || CoreAiResolver;
     const ServiceClass = options.service || CoreAiService;
+    // Registered through the persistence layer, so `database` routes them like every other model (db-agnostic phase 4);
+    // `@InjectModel()` keeps working, and each service reaches its routed store through `dataStore`
+    const aiModels = CorePersistenceModule.forFeature([
+      { model: CoreAiBudgetLimit, name: AI_BUDGET_LIMIT_MODEL, schema: AiBudgetLimitSchema },
+      { model: CoreAiConnection, name: AI_CONNECTION_MODEL, schema: AiConnectionSchema },
+      { model: CoreAiConnectionPreference, name: AI_CONNECTION_PREFERENCE_MODEL, schema: AiConnectionPreferenceSchema },
+      { model: CoreAiConversation, name: AI_CONVERSATION_MODEL, schema: AiConversationSchema },
+      { model: CoreAiInteraction, name: AI_INTERACTION_MODEL, schema: AiInteractionSchema },
+      { model: CoreAiPromptHint, name: AI_PROMPT_HINT_MODEL, schema: AiPromptHintSchema },
+      { model: CoreAiPrompt, name: AI_PROMPT_MODEL, schema: AiPromptSchema },
+      { model: CoreAiSlot, name: AI_SLOT_MODEL, schema: AiSlotSchema },
+      { model: CoreAiToolGrant, name: AI_TOOL_GRANT_MODEL, schema: AiToolGrantSchema },
+      { model: CoreAiToolPolicy, name: AI_TOOL_POLICY_MODEL, schema: AiToolPolicySchema },
+      { model: CoreAiMode, name: AI_MODE_MODEL, schema: AiModeSchema },
+      // Stored natively on MongoDB; the models give them their tables on a SQL source
+      { name: AI_MCP_OAUTH_CLIENT_MODEL, schema: AiMcpOAuthClientSchema },
+      { name: AI_MCP_OAUTH_CODE_MODEL, schema: AiMcpOAuthCodeSchema },
+      { name: AI_MCP_OAUTH_REFRESH_TOKEN_MODEL, schema: AiMcpOAuthRefreshTokenSchema },
+    ]);
     const ToolGrantServiceClass = options.toolGrantService || CoreAiToolGrantService;
     const ToolPolicyServiceClass = options.toolPolicyService || CoreAiToolPolicyService;
     const ModeServiceClass = options.modeService || CoreAiModeService;
@@ -224,23 +253,9 @@ export class CoreAiModule {
         CoreAiModeService,
         CoreAiMcpClientService,
         LlmProviderFactory,
-        MongooseModule,
-      ],
-      imports: [
-        MongooseModule.forFeature([
-          { name: AI_BUDGET_LIMIT_MODEL, schema: AiBudgetLimitSchema },
-          { name: AI_CONNECTION_MODEL, schema: AiConnectionSchema },
-          { name: AI_CONNECTION_PREFERENCE_MODEL, schema: AiConnectionPreferenceSchema },
-          { name: AI_CONVERSATION_MODEL, schema: AiConversationSchema },
-          { name: AI_INTERACTION_MODEL, schema: AiInteractionSchema },
-          { name: AI_PROMPT_HINT_MODEL, schema: AiPromptHintSchema },
-          { name: AI_PROMPT_MODEL, schema: AiPromptSchema },
-          { name: AI_SLOT_MODEL, schema: AiSlotSchema },
-          { name: AI_TOOL_GRANT_MODEL, schema: AiToolGrantSchema },
-          { name: AI_TOOL_POLICY_MODEL, schema: AiToolPolicySchema },
-          { name: AI_MODE_MODEL, schema: AiModeSchema },
-        ]),
+        aiModels,
       ],
+      imports: [aiModels],
       module: CoreAiModule,
       providers: [
         AiCryptoService,

package/src/core/modules/ai/models/core-ai-budget-limit.model.ts

@@ -1,6 +1,6 @@
 import { ObjectType } from '@nestjs/graphql';
-import { Schema as MongooseSchema, SchemaFactory } from '@nestjs/mongoose';
-import { Document } from 'mongoose';
+import { Schema as MongooseSchema, SchemaFactory } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Document } from 'mongoose';
 
 import { Restricted } from '../../../common/decorators/restricted.decorator.js';
 import { UnifiedField } from '../../../common/decorators/unified-field.decorator.js';

package/src/core/modules/ai/models/core-ai-connection-preference.model.ts

@@ -1,6 +1,6 @@
 import { ObjectType } from '@nestjs/graphql';
-import { Schema as MongooseSchema, SchemaFactory } from '@nestjs/mongoose';
-import { Document } from 'mongoose';
+import { Schema as MongooseSchema, SchemaFactory } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Document } from 'mongoose';
 
 import { Restricted } from '../../../common/decorators/restricted.decorator.js';
 import { UnifiedField } from '../../../common/decorators/unified-field.decorator.js';

package/src/core/modules/ai/models/core-ai-connection.model.ts

@@ -1,6 +1,6 @@
 import { ObjectType } from '@nestjs/graphql';
-import { Prop, Schema as MongooseSchema, SchemaFactory } from '@nestjs/mongoose';
-import { Document } from 'mongoose';
+import { Prop, Schema as MongooseSchema, SchemaFactory } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Document } from 'mongoose';
 
 import { Restricted } from '../../../common/decorators/restricted.decorator.js';
 import { UnifiedField } from '../../../common/decorators/unified-field.decorator.js';

package/src/core/modules/ai/models/core-ai-conversation.model.ts

@@ -1,6 +1,10 @@
 import { ObjectType } from '@nestjs/graphql';
-import { Schema as MongooseSchema, SchemaFactory } from '@nestjs/mongoose';
-import { Document, Schema } from 'mongoose';
+import {
+  objectIdType,
+  Schema as MongooseSchema,
+  SchemaFactory,
+} from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Document } from 'mongoose';
 
 import { Restricted } from '../../../common/decorators/restricted.decorator.js';
 import { UnifiedField } from '../../../common/decorators/unified-field.decorator.js';
@@ -40,7 +44,7 @@ export class CoreAiConversation extends CorePersistenceModel {
   @UnifiedField({
     description: 'Id of the user who created the conversation',
     isOptional: true,
-    mongoose: { index: true, ref: 'User', type: Schema.Types.ObjectId },
+    mongoose: { index: true, ref: 'User', type: objectIdType() },
     roles: RoleEnum.S_USER,
     type: () => String,
   })
@@ -69,7 +73,7 @@ export class CoreAiConversation extends CorePersistenceModel {
     description: 'Conversation messages',
     isArray: true,
     isOptional: true,
-    mongoose: { default: [], type: [Schema.Types.Mixed] },
+    mongoose: { default: [], type: [Object] },
     roles: RoleEnum.S_USER,
     type: () => CoreAiMessage,
   })

package/src/core/modules/ai/models/core-ai-interaction.model.ts

@@ -1,6 +1,6 @@
 import { ObjectType } from '@nestjs/graphql';
-import { Schema as MongooseSchema, SchemaFactory } from '@nestjs/mongoose';
-import { Document, Schema } from 'mongoose';
+import { Schema as MongooseSchema, SchemaFactory } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Document } from 'mongoose';
 
 import { Restricted } from '../../../common/decorators/restricted.decorator.js';
 import { UnifiedField } from '../../../common/decorators/unified-field.decorator.js';
@@ -26,7 +26,7 @@ export class CoreAiInteraction extends CorePersistenceModel {
   @UnifiedField({
     description: 'Tool actions executed during the run',
     isOptional: true,
-    mongoose: { type: Schema.Types.Mixed },
+    mongoose: { type: Object },
     roles: RoleEnum.ADMIN,
     type: () => JSON,
   })

package/src/core/modules/ai/models/core-ai-mcp-oauth.models.ts

@@ -0,0 +1,46 @@
+import { createSchema } from '../../../common/persistence/mongo/mongoose-compat.js';
+
+import { AI_MCP_OAUTH_CODE_TTL_SECONDS } from '../core-ai.constants.js';
+
+/**
+ * The MCP OAuth artifacts as models — what gives them their tables on a SQL source (db-agnostic phase 4).
+ *
+ * On MongoDB `CoreAiMcpOAuthService` keeps reading and writing the three collections natively, as before: these
+ * schemas only register the names, with `autoIndex: false` so Mongoose never creates an index next to the TTL and
+ * unique indexes the service ensures itself. On a SQL source the framework creates the tables from them, the service
+ * goes through their stores, and the declared TTL makes the expiry sweep remove old codes.
+ */
+
+/** A registered client: the client object the MCP SDK hands over, kept whole, plus its id for the lookup */
+export const AiMcpOAuthClientSchema = createSchema(
+  {
+    client: { type: Object },
+    client_id: { index: { unique: true }, type: String },
+    createdAt: { type: Date },
+  },
+  { autoIndex: false, collection: 'aiMcpOAuthClients', versionKey: false },
+);
+
+/** An authorization code with its PKCE challenge; valid for `AI_MCP_OAUTH_CODE_TTL_SECONDS` after `createdAt` */
+export const AiMcpOAuthCodeSchema = createSchema(
+  {
+    clientId: { type: String },
+    code: { index: { unique: true }, type: String },
+    codeChallenge: { type: String },
+    // The TTL: MongoDB gets its index from the service, a SQL table its sweep (`CorePersistenceLifecycle`)
+    createdAt: { expires: AI_MCP_OAUTH_CODE_TTL_SECONDS, type: Date },
+    userId: { type: String },
+  },
+  { autoIndex: false, collection: 'aiMcpOAuthCodes', versionKey: false },
+);
+
+/** A refresh token, bound to the client it was issued to */
+export const AiMcpOAuthRefreshTokenSchema = createSchema(
+  {
+    clientId: { type: String },
+    createdAt: { type: Date },
+    token: { index: { unique: true }, type: String },
+    userId: { type: String },
+  },
+  { autoIndex: false, collection: 'aiMcpOAuthRefreshTokens', versionKey: false },
+);

package/src/core/modules/ai/models/core-ai-mode.model.ts

@@ -1,6 +1,6 @@
 import { ObjectType } from '@nestjs/graphql';
-import { Schema as MongooseSchema, SchemaFactory } from '@nestjs/mongoose';
-import { Document } from 'mongoose';
+import { Schema as MongooseSchema, SchemaFactory } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Document } from 'mongoose';
 
 import { Restricted } from '../../../common/decorators/restricted.decorator.js';
 import { UnifiedField } from '../../../common/decorators/unified-field.decorator.js';

package/src/core/modules/ai/models/core-ai-prompt-hint.model.ts

@@ -1,6 +1,6 @@
 import { ObjectType } from '@nestjs/graphql';
-import { Schema as MongooseSchema, SchemaFactory } from '@nestjs/mongoose';
-import { Document } from 'mongoose';
+import { Schema as MongooseSchema, SchemaFactory } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Document } from 'mongoose';
 
 import { Restricted } from '../../../common/decorators/restricted.decorator.js';
 import { UnifiedField } from '../../../common/decorators/unified-field.decorator.js';

package/src/core/modules/ai/models/core-ai-prompt.model.ts

@@ -1,6 +1,6 @@
 import { ObjectType } from '@nestjs/graphql';
-import { Schema as MongooseSchema, SchemaFactory } from '@nestjs/mongoose';
-import { Document } from 'mongoose';
+import { Schema as MongooseSchema, SchemaFactory } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Document } from 'mongoose';
 
 import { Restricted } from '../../../common/decorators/restricted.decorator.js';
 import { UnifiedField } from '../../../common/decorators/unified-field.decorator.js';

package/src/core/modules/ai/models/core-ai-slot.model.ts

@@ -1,6 +1,6 @@
 import { ObjectType } from '@nestjs/graphql';
-import { Schema as MongooseSchema, SchemaFactory } from '@nestjs/mongoose';
-import { Document } from 'mongoose';
+import { Schema as MongooseSchema, SchemaFactory } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Document } from 'mongoose';
 
 import { Restricted } from '../../../common/decorators/restricted.decorator.js';
 import { UnifiedField } from '../../../common/decorators/unified-field.decorator.js';

package/src/core/modules/ai/models/core-ai-tool-grant.model.ts

@@ -1,6 +1,6 @@
 import { ObjectType } from '@nestjs/graphql';
-import { Schema as MongooseSchema, SchemaFactory } from '@nestjs/mongoose';
-import { Document } from 'mongoose';
+import { Schema as MongooseSchema, SchemaFactory } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Document } from 'mongoose';
 
 import { Restricted } from '../../../common/decorators/restricted.decorator.js';
 import { UnifiedField } from '../../../common/decorators/unified-field.decorator.js';

package/src/core/modules/ai/models/core-ai-tool-policy.model.ts

@@ -1,6 +1,6 @@
 import { Field, ObjectType } from '@nestjs/graphql';
-import { Schema as MongooseSchema, SchemaFactory } from '@nestjs/mongoose';
-import { Document } from 'mongoose';
+import { Schema as MongooseSchema, SchemaFactory } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Document } from 'mongoose';
 
 import { Restricted } from '../../../common/decorators/restricted.decorator.js';
 import { UnifiedField } from '../../../common/decorators/unified-field.decorator.js';

package/src/core/modules/ai/services/core-ai-budget.service.ts

@@ -1,9 +1,12 @@
 import { HttpException, HttpStatus, Inject, Injectable } from '@nestjs/common';
-import { InjectConnection, InjectModel } from '@nestjs/mongoose';
-import { type Connection, Model } from 'mongoose';
+import { InjectConnection, InjectModel } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Connection, Model } from 'mongoose';
 
+import type { DataStore } from '../../../common/persistence/data-store.interface.js';
+import { getModelStore } from '../../../common/persistence/persistence-runtime.registry.js';
 import { ConfigService } from '../../../common/services/config.service.js';
 import { CrudService } from '../../../common/services/crud.service.js';
+import { RequestContext } from '../../../common/services/request-context.service.js';
 import type { CoreModelConstructor } from '../../../common/types/core-model-constructor.type.js';
 import { ErrorCode } from '../../error-code/index.js';
 import { CoreAiBudgetLimitCreateInput } from '../inputs/core-ai-budget-limit-create.input.js';
@@ -11,7 +14,7 @@ import { CoreAiBudgetLimitInput } from '../inputs/core-ai-budget-limit.input.js'
 import { AiBudgetLimitDocument, CoreAiBudgetLimit } from '../models/core-ai-budget-limit.model.js';
 import { CoreAiBudgetSummary, CoreAiUsageInfo, CoreAiUsageScope } from '../models/core-ai-usage-info.model.js';
 
-import { AI_BUDGET_LIMIT_CLASS, AI_BUDGET_LIMIT_MODEL } from '../core-ai.constants.js';
+import { AI_BUDGET_LIMIT_CLASS, AI_BUDGET_LIMIT_MODEL, AI_INTERACTION_MODEL } from '../core-ai.constants.js';
 
 /**
  * @deprecated Import from `../core-ai.constants` instead. Re-exported only so existing deep imports
@@ -68,7 +71,7 @@ export class CoreAiBudgetService extends CrudService<
    * Resolve the effective limit for a scope: persisted override, else config default.
    */
   async resolveLimit(scope: 'tenant' | 'user', refId: string): Promise<ResolvedAiBudgetLimit> {
-    const override = await this.mainDbModel.findOne({ refId, scope }).lean().exec();
+    const override = (await this.dataStore.findOne({ refId, scope })) as Partial<CoreAiBudgetLimit> | null;
     const cfg = ConfigService.get<{
       period?: string;
       tenant?: { maxPrompts?: number; maxTokens?: number };
@@ -96,6 +99,15 @@ export class CoreAiBudgetService extends CrudService<
     if (filter.tenantId) {
       match.tenantId = filter.tenantId;
     }
+    // On a SQL source the interactions live in their table, summed by the database across tenants (the explicit
+    // user/tenant filter decides), as the MongoDB aggregation below does
+    const interactions = this.interactionsSqlStore();
+    if (interactions) {
+      const { count, sum } = await RequestContext.runWithBypassTenantGuard(() =>
+        interactions.sum('totalTokens', match),
+      );
+      return { resetAt: this.nextReset(period), usedPrompts: count, usedTokens: sum };
+    }
     const db = this.connection.db;
     if (!db) {
       return { resetAt: this.nextReset(period), usedPrompts: 0, usedTokens: 0 };
@@ -113,6 +125,12 @@ export class CoreAiBudgetService extends CrudService<
     return { resetAt: this.nextReset(period), usedPrompts: agg?.usedPrompts ?? 0, usedTokens: agg?.usedTokens ?? 0 };
   }
 
+  /** The store of the interactions when they are routed to a SQL source; `undefined` keeps the native aggregation */
+  protected interactionsSqlStore(): DataStore<any> | undefined {
+    const store = getModelStore(this.connection, AI_INTERACTION_MODEL);
+    return store && !store.capabilities.mongoosePopulate ? store : undefined;
+  }
+
   /**
    * Block the run with HTTP 429 if the user OR tenant has hit a finite limit.
    *

package/src/core/modules/ai/services/core-ai-connection-preference.service.ts

@@ -1,7 +1,8 @@
 import { Inject, Injectable } from '@nestjs/common';
-import { InjectModel } from '@nestjs/mongoose';
-import { Model } from 'mongoose';
+import { InjectModel } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Model } from 'mongoose';
 
+import { upsertRecord } from '../../../common/persistence/store-upsert.js';
 import { CrudService } from '../../../common/services/crud.service.js';
 import type { CoreModelConstructor } from '../../../common/types/core-model-constructor.type.js';
 import { CoreAiConnectionPreferenceInput } from '../inputs/core-ai-connection-preference.input.js';
@@ -45,7 +46,7 @@ export class CoreAiConnectionPreferenceService extends CrudService<
     scope: 'tenant' | 'user',
     refId: string,
   ): Promise<{ connectionId: string; enforced?: boolean } | null> {
-    const doc = await this.mainDbModel.findOne({ refId, scope }).lean().exec();
+    const doc = (await this.dataStore.findOne({ refId, scope })) as Partial<CoreAiConnectionPreference> | null;
     return doc ? { connectionId: doc.connectionId, enforced: doc.enforced } : null;
   }
 
@@ -54,8 +55,7 @@ export class CoreAiConnectionPreferenceService extends CrudService<
    * removed, to avoid dangling tenant/user preferences). Returns the deleted count.
    */
   async deleteByConnectionId(connectionId: string): Promise<number> {
-    const result = await this.mainDbModel.deleteMany({ connectionId }).exec();
-    return result.deletedCount ?? 0;
+    return this.dataStore.deleteMany({ connectionId });
   }
 
   /**
@@ -69,14 +69,7 @@ export class CoreAiConnectionPreferenceService extends CrudService<
     connectionId: string,
     enforced = false,
   ): Promise<CoreAiConnectionPreference> {
-    const doc = await this.mainDbModel
-      .findOneAndUpdate(
-        { refId, scope },
-        { $set: { connectionId, enforced } },
-        { returnDocument: 'after', upsert: true },
-      )
-      .lean()
-      .exec();
+    const doc = await upsertRecord(this.dataStore, { refId, scope }, { connectionId, enforced });
     return this.mainModelConstructor.map(doc);
   }
 }

package/src/core/modules/ai/services/core-ai-connection.service.ts

@@ -7,8 +7,8 @@ import {
   Optional,
   ServiceUnavailableException,
 } from '@nestjs/common';
-import { InjectModel } from '@nestjs/mongoose';
-import { Model } from 'mongoose';
+import { InjectModel } from '../../../common/persistence/mongo/mongoose-compat.js';
+import type { Model } from 'mongoose';
 
 import { ServiceOptions } from '../../../common/interfaces/service-options.interface.js';
 import { CrudService } from '../../../common/services/crud.service.js';
@@ -39,7 +39,7 @@ export { AI_CONNECTION_CLASS, AI_CONNECTION_MODEL } from '../core-ai.constants.j
  * use the bulk `find()` result directly (no per-connection re-read / N+1).
  */
 type ResolvableConnectionDoc = {
-  _id: unknown;
+  id: string;
   apiKeyEncrypted?: string;
   apiKeyEnv?: string;
   baseUrl: string;
@@ -121,14 +121,14 @@ export class CoreAiConnectionService
       );
       return;
     }
-    const count = await this.mainDbModel.countDocuments().exec();
+    const count = await this.dataStore.count();
     if (count > 0) {
       return;
     }
     const { apiKey, ...rest } = seed;
-    // Direct model access: system-internal seeding without a user context, and we
+    // Direct store access: system-internal seeding without a user context, and we
     // must set the internal `apiKeyEncrypted` field which CrudService does not expose.
-    await this.mainDbModel.create({
+    await this.dataStore.insert({
       ...rest,
       apiKeyEncrypted: apiKey ? this.aiCryptoService.encrypt(apiKey) : undefined,
       isDefault: true,
@@ -145,16 +145,16 @@ export class CoreAiConnectionService
    */
   protected async assertStoredKeysDecryptable(): Promise<void> {
     try {
-      const docs = await this.mainDbModel
-        .find({ apiKeyEncrypted: { $exists: true, $ne: '' } }, { apiKeyEncrypted: 1, name: 1 })
-        .lean()
-        .exec();
+      const docs = await this.dataStore.find(
+        { apiKeyEncrypted: { $exists: true, $ne: '' } },
+        { select: ['apiKeyEncrypted', 'name'] },
+      );
       const broken: string[] = [];
-      for (const doc of docs as { _id: unknown; apiKeyEncrypted?: string; name?: string }[]) {
+      for (const doc of docs as { apiKeyEncrypted?: string; id: string; name?: string }[]) {
         try {
           this.aiCryptoService.decrypt(doc.apiKeyEncrypted ?? '');
         } catch {
-          broken.push(doc.name || String(doc._id));
+          broken.push(doc.name || doc.id);
         }
       }
       if (broken.length) {
@@ -205,10 +205,7 @@ export class CoreAiConnectionService
     try {
       // Single read (no per-connection re-resolve): the full docs carry everything the
       // provider factory needs, so there is no N+1 findById per connection.
-      const docs = (await this.mainDbModel
-        .find({ enabled: { $ne: false } })
-        .lean()
-        .exec()) as unknown as ResolvableConnectionDoc[];
+      const docs = (await this.dataStore.find({ enabled: { $ne: false } })) as unknown as ResolvableConnectionDoc[];
       for (const doc of docs) {
         // Only a connection that DECLARES a capability can drift; undefined flags are
         // auto-detected on first use, so there is nothing to reconcile here.
@@ -227,7 +224,7 @@ export class CoreAiConnectionService
             defaultTemperature: doc.defaultTemperature,
             defaultUserMaxPeriod: doc.defaultUserMaxPeriod,
             defaultUserMaxTokens: doc.defaultUserMaxTokens,
-            id: String(doc._id),
+            id: String(doc.id),
             model: doc.model,
             name: doc.name,
             providerType: doc.providerType || 'openai-compatible',
@@ -266,7 +263,7 @@ export class CoreAiConnectionService
         }
         if (drift.length) {
           this.logger.warn(
-            `AI connection "${doc.name || String(doc._id)}" capability drift: ${drift.join('; ')}. ` +
+            `AI connection "${doc.name || String(doc.id)}" capability drift: ${drift.join('; ')}. ` +
               `The declared value is authoritative and was NOT changed — correct it in the admin UI, or clear it to ` +
               `re-enable auto-detection, so the assistant uses the endpoint's real capabilities.`,
           );
@@ -399,7 +396,7 @@ export class CoreAiConnectionService
       }
     }
     if (Object.keys(update).length) {
-      await this.mainDbModel.findByIdAndUpdate(connectionId, { $set: update }).exec();
+      await this.dataStore.updateById(connectionId, { $set: update });
       this.logger.log(`Auto-detected AI connection metadata for "${resolved.name}": ${JSON.stringify(update)}`);
     }
     return resolved;
@@ -420,16 +417,14 @@ export class CoreAiConnectionService
       tenantIds?: string[];
     }[]
   > {
-    const docs = await this.mainDbModel
-      .find({ enabled: { $ne: false } })
-      .select('enforced enforcedTenantIds isDefault model name tenantIds')
-      .sort({ createdAt: 1 })
-      .lean()
-      .exec();
+    const docs = (await this.dataStore.find(
+      { enabled: { $ne: false } },
+      { select: ['enforced', 'enforcedTenantIds', 'isDefault', 'model', 'name', 'tenantIds'], sort: { createdAt: 1 } },
+    )) as (Partial<CoreAiConnection> & { id: string })[];
     return docs.map((doc) => ({
       enforced: doc.enforced,
       enforcedTenantIds: doc.enforcedTenantIds,
-      id: String(doc._id),
+      id: String(doc.id),
       isDefault: doc.isDefault,
       model: doc.model,
       name: doc.name,
@@ -444,21 +439,18 @@ export class CoreAiConnectionService
    * @param connectionId Specific connection, or undefined for the default / first enabled one.
    */
   async resolve(connectionId?: string): Promise<ResolvedAiConnection> {
-    // Direct lean reads: system-internal resolution that must access the
+    // Direct store reads: system-internal resolution that must access the
     // encrypted key field, which CrudService strips from outputs.
-    let doc = connectionId
-      ? await this.mainDbModel.findById(connectionId).lean().exec()
-      : await this.mainDbModel
-          .findOne({ enabled: { $ne: false }, isDefault: true })
-          .lean()
-          .exec();
+    let doc = (
+      connectionId
+        ? await this.dataStore.findById(connectionId)
+        : await this.dataStore.findOne({ enabled: { $ne: false }, isDefault: true })
+    ) as (CoreAiConnection & { id: string }) | null;
 
     if (!doc && !connectionId) {
-      doc = await this.mainDbModel
-        .findOne({ enabled: { $ne: false } })
-        .sort({ createdAt: 1 })
-        .lean()
-        .exec();
+      doc = (await this.dataStore.findOne({ enabled: { $ne: false } }, { sort: { createdAt: 1 } })) as
+        | (CoreAiConnection & { id: string })
+        | null;
     }
 
     if (!doc) {
@@ -476,7 +468,7 @@ export class CoreAiConnectionService
       defaultTemperature: doc.defaultTemperature,
       defaultUserMaxPeriod: doc.defaultUserMaxPeriod,
       defaultUserMaxTokens: doc.defaultUserMaxTokens,
-      id: String(doc._id),
+      id: String(doc.id),
       model: doc.model,
       name: doc.name,
       providerType: doc.providerType || 'openai-compatible',
@@ -490,11 +482,9 @@ export class CoreAiConnectionService
    */
   protected async applyApiKey(id: string, apiKey: string): Promise<void> {
     if (apiKey === '') {
-      await this.mainDbModel.findByIdAndUpdate(id, { $unset: { apiKeyEncrypted: 1 } }).exec();
+      await this.dataStore.updateById(id, { $unset: { apiKeyEncrypted: 1 } });
     } else {
-      await this.mainDbModel
-        .findByIdAndUpdate(id, { $set: { apiKeyEncrypted: this.aiCryptoService.encrypt(apiKey) } })
-        .exec();
+      await this.dataStore.updateById(id, { $set: { apiKeyEncrypted: this.aiCryptoService.encrypt(apiKey) } });
     }
   }
 
@@ -516,6 +506,6 @@ export class CoreAiConnectionService
    * Ensure only one connection is marked as default.
    */
   protected async unsetOtherDefaults(id: string): Promise<void> {
-    await this.mainDbModel.updateMany({ _id: { $ne: id } }, { $set: { isDefault: false } }).exec();
+    await this.dataStore.updateMany({ id: { $ne: id } }, { $set: { isDefault: false } });
   }
 }