Merge pull request 'feat(ai-chat): agent roles (admin persona + optional model)' (#11) from feat/ai-agent-roles into develop

apps/server/src/integrations/ai/ai-not-configured.exception.ts

@@ -5,7 +5,7 @@ import { ServiceUnavailableException } from '@nestjs/common';
  * driver / chat model / API key). Maps to HTTP 503 (§6.2/§6.4).
  */
 export class AiNotConfiguredException extends ServiceUnavailableException {
-  constructor() {
-    super('AI provider not configured');
+  constructor(message = 'AI provider not configured') {
+    super(message);
   }
 }

apps/server/src/integrations/ai/ai.service.spec.ts

@@ -0,0 +1,174 @@
+import { AiService } from './ai.service';
+import { AiNotConfiguredException } from './ai-not-configured.exception';
+
+/**
+ * Unit test for the role model-override 503 path of AiService.getChatModel.
+ *
+ * AiService's constructor body is trivial (it only stores its deps), so it can
+ * be unit-constructed with stubbed collaborators — no Nest module graph, which
+ * the src-rooted jest setup cannot fully resolve for the heavier specs. We stub:
+ *  - aiSettings.resolve  -> a workspace configured for openai (so cfg.driver is
+ *    set and we pass the first guard),
+ *  - aiProviderCredentialsRepo.find -> undefined (the override driver has NO
+ *    configured credentials),
+ *  - secretBox -> unused on this path (no creds to decrypt).
+ *
+ * With a role override pointing at a DIFFERENT driver ('gemini') that has no
+ * creds, getChatModel must throw AiNotConfiguredException (503) and the message
+ * must name the override driver (and the role) so an admin can fix it.
+ */
+describe('AiService.getChatModel role model override', () => {
+  function makeService(opts: {
+    workspaceDriver: string;
+    credsApiKeyEnc?: string;
+  }) {
+    const aiSettings = {
+      resolve: jest.fn().mockResolvedValue({
+        driver: opts.workspaceDriver,
+        chatModel: 'gpt-4o-mini',
+        apiKey: 'workspace-key',
+        baseUrl: undefined,
+      }),
+    };
+    const aiProviderCredentialsRepo = {
+      find: jest.fn().mockResolvedValue(
+        opts.credsApiKeyEnc ? { apiKeyEnc: opts.credsApiKeyEnc } : undefined,
+      ),
+    };
+    const secretBox = {
+      decryptSecret: jest.fn().mockReturnValue('decrypted'),
+    };
+    const service = new AiService(
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      aiSettings as any,
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      aiProviderCredentialsRepo as any,
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      secretBox as any,
+    );
+    return { service, aiSettings, aiProviderCredentialsRepo, secretBox };
+  }
+
+  it('throws AiNotConfiguredException (503) naming the override driver when its creds are missing', async () => {
+    const { service, aiProviderCredentialsRepo } = makeService({
+      workspaceDriver: 'openai',
+    });
+
+    await expect(
+      service.getChatModel('ws-1', {
+        driver: 'gemini',
+        chatModel: 'gemini-2.0-flash',
+        roleName: 'Researcher',
+      }),
+    ).rejects.toBeInstanceOf(AiNotConfiguredException);
+
+    // Re-run to assert the message names the driver (and role) for the admin.
+    await service
+      .getChatModel('ws-1', {
+        driver: 'gemini',
+        chatModel: 'gemini-2.0-flash',
+        roleName: 'Researcher',
+      })
+      .then(
+        () => {
+          throw new Error('expected getChatModel to throw');
+        },
+        (err: unknown) => {
+          expect(err).toBeInstanceOf(AiNotConfiguredException);
+          const message = (err as AiNotConfiguredException).message;
+          expect(message).toContain('gemini');
+          expect(message).toContain('Researcher');
+        },
+      );
+
+    // The override driver's creds were looked up for the right driver.
+    expect(aiProviderCredentialsRepo.find).toHaveBeenCalledWith('ws-1', 'gemini');
+  });
+
+  it('cross-driver override with creds present: resolves without throwing, using the OVERRIDE driver creds', async () => {
+    // Workspace driver is openai; the role overrides to gemini, which HAS creds.
+    const { service, aiProviderCredentialsRepo, secretBox } = makeService({
+      workspaceDriver: 'openai',
+      credsApiKeyEnc: 'enc-gemini-key',
+    });
+
+    const model = await service.getChatModel('ws-1', {
+      driver: 'gemini',
+      chatModel: 'gemini-2.0-flash',
+      roleName: 'Researcher',
+    });
+
+    // A real LanguageModel was built (no 503).
+    expect(model).toBeDefined();
+    // Creds were fetched for the OVERRIDE driver, then decrypted.
+    expect(aiProviderCredentialsRepo.find).toHaveBeenCalledWith('ws-1', 'gemini');
+    expect(secretBox.decryptSecret).toHaveBeenCalledWith('enc-gemini-key');
+  });
+
+  it('cross-driver override to ollama (workspace driver != ollama): throws 503, does NOT silently reuse the workspace baseUrl', async () => {
+    // Workspace driver is openai with a configured (gateway) baseUrl. A role that
+    // overrides to ollama has no dedicated ollama endpoint, so pointing the
+    // ollama client at the workspace's openai baseUrl would be wrong — it must
+    // fail explicitly instead.
+    const aiSettings = {
+      resolve: jest.fn().mockResolvedValue({
+        driver: 'openai',
+        chatModel: 'gpt-4o-mini',
+        apiKey: 'workspace-key',
+        baseUrl: 'https://openrouter.example/v1',
+      }),
+    };
+    const aiProviderCredentialsRepo = { find: jest.fn() };
+    const secretBox = { decryptSecret: jest.fn() };
+    const service = new AiService(
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      aiSettings as any,
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      aiProviderCredentialsRepo as any,
+      // eslint-disable-next-line @typescript-eslint/no-explicit-any
+      secretBox as any,
+    );
+
+    await service
+      .getChatModel('ws-1', {
+        driver: 'ollama',
+        chatModel: 'llama3',
+        roleName: 'Local',
+      })
+      .then(
+        () => {
+          throw new Error('expected getChatModel to throw');
+        },
+        (err: unknown) => {
+          expect(err).toBeInstanceOf(AiNotConfiguredException);
+          const message = (err as AiNotConfiguredException).message;
+          // Names the role and the workspace driver, and mentions ollama.
+          expect(message).toContain('ollama');
+          expect(message).toContain('openai');
+          expect(message).toContain('Local');
+          // Must NOT leak / reuse the workspace gateway baseUrl in the path.
+          expect(message).not.toContain('openrouter.example');
+        },
+      );
+
+    // No ollama creds lookup happens (ollama needs no key); we fail before that.
+    expect(aiProviderCredentialsRepo.find).not.toHaveBeenCalled();
+  });
+
+  it('chatModel-only override (no driver): reuses the workspace driver+creds, no creds lookup/decrypt', async () => {
+    // No override.driver => the workspace openai driver + its apiKey are reused;
+    // ai_provider_credentials must NOT be queried and nothing is decrypted.
+    const { service, aiProviderCredentialsRepo, secretBox } = makeService({
+      workspaceDriver: 'openai',
+    });
+
+    const model = await service.getChatModel('ws-1', {
+      chatModel: 'gpt-4o',
+      roleName: 'Writer',
+    });
+
+    expect(model).toBeDefined();
+    expect(aiProviderCredentialsRepo.find).not.toHaveBeenCalled();
+    expect(secretBox.decryptSecret).not.toHaveBeenCalled();
+  });
+});

apps/server/src/integrations/ai/ai.service.ts

@@ -14,6 +14,22 @@ import { AiNotConfiguredException } from './ai-not-configured.exception';
 import { AiEmbeddingNotConfiguredException } from './ai-embedding-not-configured.exception';
 import { AiSttNotConfiguredException } from './ai-stt-not-configured.exception';
 import { describeProviderError } from './ai-error.util';
+import { AiProviderCredentialsRepo } from '@docmost/db/repos/ai-chat/ai-provider-credentials.repo';
+import { SecretBoxService } from '../crypto/secret-box';
+import { AiDriver } from './ai.types';
+
+/**
+ * Optional chat-model override carried by an agent role (`ai_agent_roles.
+ * model_config`). `chatModel` swaps the model id; `driver` (optional) switches
+ * the whole provider, in which case its creds come from `ai_provider_credentials`
+ * for that driver. `roleName` is only used to produce a clear 503 message when
+ * the chosen driver is not configured.
+ */
+export interface ChatModelOverride {
+  driver?: AiDriver;
+  chatModel?: string;
+  roleName?: string;
+}
 
 /**
  * Builds AI SDK language models from per-workspace config and runs cheap
@@ -27,23 +43,91 @@ import { describeProviderError } from './ai-error.util';
 export class AiService {
   private readonly logger = new Logger(AiService.name);
 
-  constructor(private readonly aiSettings: AiSettingsService) {}
+  constructor(
+    private readonly aiSettings: AiSettingsService,
+    private readonly aiProviderCredentialsRepo: AiProviderCredentialsRepo,
+    private readonly secretBox: SecretBoxService,
+  ) {}
 
   /**
    * Resolve the workspace config and build the chat language model.
    * Throws AiNotConfiguredException (→ 503) when the config is incomplete.
+   *
+   * `override` (from an agent role's `model_config`) optionally swaps the model
+   * id and/or the whole provider:
+   *  - `override.chatModel` replaces the workspace chat model id;
+   *  - `override.driver` (when it differs from the workspace driver) switches the
+   *    provider, pulling that driver's creds from `ai_provider_credentials`. When
+   *    those creds are missing the call throws a 503 naming the role's driver — a
+   *    deliberate, explicit failure rather than a silent fallback. Resolved
+   *    BEFORE the stream starts so the 503 surfaces as clean JSON.
    */
-  async getChatModel(workspaceId: string): Promise<LanguageModel> {
+  async getChatModel(
+    workspaceId: string,
+    override?: ChatModelOverride,
+  ): Promise<LanguageModel> {
     const cfg = await this.aiSettings.resolve(workspaceId);
-    if (
-      !cfg?.driver ||
-      !cfg?.chatModel ||
-      (cfg.driver !== 'ollama' && !cfg.apiKey)
-    ) {
+    if (!cfg?.driver) {
       throw new AiNotConfiguredException();
     }
 
-    switch (cfg.driver) {
+    // Determine the effective driver + model + creds, applying the override.
+    const overrideDriver = override?.driver;
+    const driver: AiDriver = overrideDriver ?? cfg.driver;
+    const chatModel = override?.chatModel?.trim() || cfg.chatModel;
+
+    let apiKey = cfg.apiKey;
+    let baseUrl = cfg.baseUrl;
+
+    // A driver override that differs from the workspace driver needs that
+    // driver's own creds (the workspace driver's key would be wrong/absent).
+    if (overrideDriver && overrideDriver !== cfg.driver) {
+      if (overrideDriver === 'ollama') {
+        // Cross-driver override to ollama: the workspace driver is NOT ollama, so
+        // there is no configured ollama endpoint. `cfg.baseUrl` belongs to the
+        // workspace driver (e.g. an OpenAI/OpenRouter gateway) and pointing the
+        // ollama client at it would silently send requests to the wrong server.
+        // Fail explicitly (503) — a dedicated per-driver ollama endpoint is not
+        // supported yet. The same-driver ollama case (handled outside this block)
+        // legitimately reuses the workspace's ollama endpoint and is unaffected.
+        const who = override?.roleName ? ` for role "${override.roleName}"` : '';
+        throw new AiNotConfiguredException(
+          `An ollama model override${who} requires a dedicated ollama endpoint, ` +
+            `which is not supported when the workspace driver is "${cfg.driver}". ` +
+            `Set the role's driver to "${cfg.driver}" or switch the workspace ` +
+            `to ollama.`,
+        );
+      } else {
+        const creds = await this.aiProviderCredentialsRepo.find(
+          workspaceId,
+          overrideDriver,
+        );
+        apiKey = creds?.apiKeyEnc
+          ? this.secretBox.decryptSecret(creds.apiKeyEnc)
+          : undefined;
+        if (!apiKey) {
+          // Explicit 503: the role chose a provider that is not set up. Name the
+          // driver (and role, when known) so the admin can fix it — no silent
+          // fallback to the workspace model (error-handling convention).
+          const who = override?.roleName ? ` for role "${override.roleName}"` : '';
+          throw new AiNotConfiguredException(
+            `The model provider "${overrideDriver}"${who} is selected but not ` +
+              `configured (no API key). Configure ${overrideDriver} in AI ` +
+              `settings or change the role's model.`,
+          );
+        }
+        // A cross-driver override does not carry the workspace baseUrl (that URL
+        // belongs to the workspace driver); use the provider default for the
+        // overridden driver.
+        baseUrl = undefined;
+      }
+    }
+
+    if (!chatModel || (driver !== 'ollama' && !apiKey)) {
+      throw new AiNotConfiguredException();
+    }
+
+    switch (driver) {
       case 'openai':
         // baseURL (when set) covers openai-compatible endpoints. Use Chat
         // Completions (/chat/completions) — the portable OpenAI-compatible
@@ -51,14 +135,12 @@ export class AiService {
         // Responses API (/responses), which OpenAI-compatible gateways
         // (OpenRouter, etc.) reject on multi-turn requests (history with
         // assistant messages) → 400.
-        return createOpenAI({ apiKey: cfg.apiKey, baseURL: cfg.baseUrl }).chat(
-          cfg.chatModel,
-        );
+        return createOpenAI({ apiKey, baseURL: baseUrl }).chat(chatModel);
       case 'gemini':
-        return createGoogleGenerativeAI({ apiKey: cfg.apiKey })(cfg.chatModel);
+        return createGoogleGenerativeAI({ apiKey })(chatModel);
       case 'ollama':
         // Ollama needs no API key.
-        return createOllama({ baseURL: cfg.baseUrl })(cfg.chatModel);
+        return createOllama({ baseURL: baseUrl })(chatModel);
       default:
         throw new AiNotConfiguredException();
     }