From e103d1231d9caba238e65e54fa8ac7aa327afeb4 Mon Sep 17 00:00:00 2001
From: Vincent Koc <vincentkoc@ieee.org>
Date: Tue, 9 Jun 2026 14:27:07 +0900
Subject: [PATCH] fix(microsoft-foundry): classify manual MAI image setup

---
 extensions/microsoft-foundry/index.test.ts | 44 ++++++++++++++++++-
 extensions/microsoft-foundry/onboard.ts    | 51 +++++++++++++++++++++-
 2 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/extensions/microsoft-foundry/index.test.ts b/extensions/microsoft-foundry/index.test.ts
index 05db6634215d..23e325d261e7 100644
--- a/extensions/microsoft-foundry/index.test.ts
+++ b/extensions/microsoft-foundry/index.test.ts
@@ -5,7 +5,11 @@ import { createTestPluginApi } from "openclaw/plugin-sdk/plugin-test-api";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { getAccessTokenResultAsync } from "./cli.js";
 import plugin from "./index.js";
-import { buildFoundryConnectionTest, isValidTenantIdentifier } from "./onboard.js";
+import {
+  buildFoundryConnectionTest,
+  isValidTenantIdentifier,
+  promptApiKeyEndpointAndModel,
+} from "./onboard.js";
 import { resetFoundryRuntimeAuthCaches } from "./runtime.js";
 import {
   buildFoundryAuthResult,
@@ -686,6 +690,44 @@ describe("microsoft-foundry plugin", () => {
     expect(requireFoundryProviderPatch(result).models[0]?.name).toBe("MAI-Image-2.5");
   });
 
+  it("classifies custom API-key MAI image deployments during manual setup", async () => {
+    const text = vi
+      .fn()
+      .mockResolvedValueOnce("https://example.services.ai.azure.com")
+      .mockResolvedValueOnce("prod-image");
+    const select = vi
+      .fn()
+      .mockResolvedValueOnce("mai-image")
+      .mockResolvedValueOnce("MAI-Image-2.5");
+    const selection = await promptApiKeyEndpointAndModel({
+      prompter: {
+        text,
+        select,
+      },
+    } as never);
+
+    const result = buildFoundryAuthResult({
+      profileId: "microsoft-foundry:default",
+      apiKey: "test-api-key",
+      endpoint: selection.endpoint,
+      modelId: selection.modelId,
+      modelNameHint: selection.modelNameHint,
+      api: selection.api,
+      authMethod: "api-key",
+    });
+
+    expect(selection).toEqual({
+      endpoint: "https://example.services.ai.azure.com",
+      modelId: "prod-image",
+      modelNameHint: "MAI-Image-2.5",
+      api: "openai-completions",
+    });
+    expect(result.configPatch?.agents?.defaults?.imageGenerationModel).toEqual({
+      primary: "microsoft-foundry/prod-image",
+    });
+    expect(result.defaultModel).toBeUndefined();
+  });
+
   it("uses discovered deployment metadata for MAI image defaults", () => {
     const result = buildFoundryAuthResult({
       profileId: "microsoft-foundry:entra",
diff --git a/extensions/microsoft-foundry/onboard.ts b/extensions/microsoft-foundry/onboard.ts
index d6b543fa00c3..5f7eb1352546 100644
--- a/extensions/microsoft-foundry/onboard.ts
+++ b/extensions/microsoft-foundry/onboard.ts
@@ -217,7 +217,12 @@ async function promptFoundryApi(
   });
 }
 
-type ManualFoundryModelFamilyChoice = "reasoning-family" | "other-chat";
+type ManualFoundryModelFamilyChoice = "reasoning-family" | "mai-image" | "other-chat";
+type ManualFoundryMaiImageModel =
+  | "MAI-Image-2.5-Flash"
+  | "MAI-Image-2.5"
+  | "MAI-Image-2e"
+  | "MAI-Image-2";
 
 async function promptFoundryModelFamily(
   ctx: ProviderAuthContext,
@@ -230,6 +235,11 @@ async function promptFoundryModelFamily(
         label: "GPT-5 series / o-series / Codex",
         hint: "Use for Azure OpenAI reasoning and Codex deployments",
       },
+      {
+        value: "mai-image",
+        label: "MAI image model",
+        hint: "Use for Microsoft MAI image deployments",
+      },
       {
         value: "other-chat",
         label: "Other chat model",
@@ -240,6 +250,37 @@ async function promptFoundryModelFamily(
   });
 }
 
+async function promptFoundryMaiImageModel(
+  ctx: ProviderAuthContext,
+): Promise<ManualFoundryMaiImageModel> {
+  return await ctx.prompter.select({
+    message: "MAI image base model",
+    options: [
+      {
+        value: "MAI-Image-2.5-Flash",
+        label: "MAI-Image-2.5-Flash",
+        hint: "Latest fast MAI image deployment",
+      },
+      {
+        value: "MAI-Image-2.5",
+        label: "MAI-Image-2.5",
+        hint: "Latest MAI image deployment",
+      },
+      {
+        value: "MAI-Image-2e",
+        label: "MAI-Image-2e",
+        hint: "Efficient MAI image deployment",
+      },
+      {
+        value: "MAI-Image-2",
+        label: "MAI-Image-2",
+        hint: "MAI image deployment",
+      },
+    ],
+    initialValue: "MAI-Image-2.5-Flash",
+  });
+}
+
 async function promptEndpointAndModelBase(
   ctx: ProviderAuthContext,
   options?: {
@@ -276,6 +317,14 @@ async function promptEndpointAndModelBase(
     })
   ).trim();
   const familyChoice = await promptFoundryModelFamily(ctx);
+  if (familyChoice === "mai-image") {
+    return {
+      endpoint,
+      modelId,
+      modelNameHint: await promptFoundryMaiImageModel(ctx),
+      api: DEFAULT_API,
+    };
+  }
   const resolvedModelName =
     familyChoice === "reasoning-family"
       ? usesFoundryResponsesByDefault(modelId) || requiresFoundryMaxCompletionTokens(modelId)