Studio: Add custom provider option to Connections (#6112)

* feat: add custom connection * Fix custom provider handling for PR #6112 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix custom provider connection test for PR #6112 --------- Co-authored-by: wasimysaid <wasimysdev@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-12 12:09:35 +01:00 · 2026-06-12 12:09:35 +01:00 · 7307fde839
commit 7307fde839
parent 72e67ae5a6
7 changed files with 229 additions and 10 deletions
--- a/studio/backend/core/inference/providers.py
+++ b/studio/backend/core/inference/providers.py
@ -249,6 +249,23 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
        # Surfaced via the frontend's CUSTOM_PROVIDER_PRESETS, not the dropdown.
        "hidden": True,
    },
+    "custom": {
+        "display_name": "Custom",
+        # User-supplied via provider_base_url.
+        "base_url": "",
+        "default_models": [],
+        "supports_streaming": True,
+        "supports_vision": True,
+        "supports_tool_calling": True,
+        "auth_header": "Authorization",
+        "auth_prefix": "Bearer ",
+        "notes": (
+            "User-supplied OpenAI-compatible server. Routed to "
+            "/v1/chat/completions; /models is optional."
+        ),
+        # Surfaced by the frontend's generic Custom option, not the dropdown.
+        "hidden": True,
+    },
    "ollama": {
        "display_name": "Ollama",
        "base_url": "http://localhost:11434/v1",
--- a/studio/backend/models/providers.py
+++ b/studio/backend/models/providers.py
@ -108,6 +108,9 @@ class ProviderTestRequest(BaseModel):
    base_url: Optional[str] = Field(
        None, description = "Custom base URL (overrides registry default)"
    )
+    model_id: Optional[str] = Field(
+        None, description = "Model ID for providers that need a chat probe"
+    )


 class ProviderTestResult(BaseModel):
--- a/studio/backend/routes/providers.py
+++ b/studio/backend/routes/providers.py
@ -190,7 +190,8 @@ async def test_provider(
    """
    Test connectivity to an external provider.

-    Makes a lightweight GET /models call to verify the API key works.
+    Makes a lightweight GET /models call to verify the API key works. Generic
+    custom endpoints use a chat-completions probe because /models is optional.
    encrypted_api_key is decrypted server-side and never stored.
    """
    info = get_provider_info(payload.provider_type)
@ -212,6 +213,14 @@ async def test_provider(
            )

    base_url = payload.base_url or info["base_url"]
+    if payload.provider_type == "custom":
+        if not base_url:
+            return ProviderTestResult(
+                success = False,
+                message = "Connection failed: Base URL is required for custom providers.",
+                models_count = None,
+            )
+
    client = ExternalProviderClient(
        provider_type = payload.provider_type,
        base_url = base_url,
@ -220,6 +229,26 @@ async def test_provider(
    )

    try:
+        if payload.provider_type == "custom":
+            model_id = (payload.model_id or "").strip()
+            if not model_id:
+                return ProviderTestResult(
+                    success = False,
+                    message = "Connection failed: add a model ID to test custom providers.",
+                    models_count = None,
+                )
+            await client.chat_completion(
+                messages = [{"role": "user", "content": "ping"}],
+                model = model_id,
+                temperature = 0.0,
+                top_p = 1.0,
+                max_tokens = 1,
+            )
+            return ProviderTestResult(
+                success = True,
+                message = "Connected successfully. Chat completions endpoint responded.",
+                models_count = None,
+            )
        if info.get("model_list_mode") == "curated":
            await client.verify_models_endpoint_lightweight()
            return ProviderTestResult(
--- a/studio/backend/tests/test_external_provider_usage_chunk.py
+++ b/studio/backend/tests/test_external_provider_usage_chunk.py
@ -144,6 +144,14 @@ def _make_openai_client() -> ExternalProviderClient:
    )


+def _make_custom_client() -> ExternalProviderClient:
+    return ExternalProviderClient(
+        provider_type = "custom",
+        base_url = "http://custom.example/v1",
+        api_key = "",
+    )
+
+
 def _anthropic_sse(events: list[dict]) -> bytes:
    chunks: list[str] = []
    for event in events:
@ -180,6 +188,137 @@ def _usage_chunks(lines: list[str]) -> list[dict]:
    return out


+def test_custom_provider_registry_is_hidden():
+    from core.inference.providers import get_provider_info, list_available_providers
+
+    info = get_provider_info("custom")
+    assert info is not None
+    assert info["hidden"] is True
+    assert "custom" not in {p["provider_type"] for p in list_available_providers()}
+
+
+def test_custom_provider_uses_chat_completions_without_auth_key(monkeypatch):
+    captured: dict = {}
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        captured["url"] = str(request.url)
+        captured["headers"] = dict(request.headers)
+        captured["body"] = json.loads(request.content.decode("utf-8"))
+        return httpx.Response(
+            200,
+            content = b'data: {"choices":[{"delta":{"content":"ok"}}]}\n\ndata: [DONE]\n\n',
+            headers = {"content-type": "text/event-stream"},
+        )
+
+    _mock_http_client(monkeypatch, handler)
+
+    async def run():
+        client = _make_custom_client()
+        lines = await _collect(
+            client.stream_chat_completion(
+                messages = [{"role": "user", "content": "ping"}],
+                model = "Qwen/Qwen3-0.6B",
+                temperature = 0.7,
+                top_p = 0.95,
+                max_tokens = 64,
+            )
+        )
+        await client.close()
+        return lines
+
+    lines = _drive(run())
+    assert captured["url"] == "http://custom.example/v1/chat/completions"
+    assert "authorization" not in {k.lower() for k in captured["headers"]}
+    assert captured["body"]["model"] == "Qwen/Qwen3-0.6B"
+    assert any("ok" in line for line in lines)
+
+
+def test_custom_provider_test_endpoint_probes_chat_completion(monkeypatch):
+    import importlib.util
+    import sys
+    from pathlib import Path
+
+    module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py"
+    spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path)
+    assert spec is not None
+    assert spec.loader is not None
+    providers_route = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = providers_route
+    spec.loader.exec_module(providers_route)
+
+    captured: dict = {}
+
+    class _FakeClient:
+        def __init__(self, **kwargs):
+            captured["init"] = kwargs
+
+        async def chat_completion(self, **kwargs):
+            captured["chat_completion"] = kwargs
+            return {"choices": [{"message": {"content": "ok"}}]}
+
+        async def list_models(self):
+            raise AssertionError("custom provider test must not call /models")
+
+        async def close(self):
+            captured["closed"] = True
+
+    monkeypatch.setattr(providers_route, "ExternalProviderClient", _FakeClient)
+
+    async def run():
+        return await providers_route.test_provider(
+            providers_route.ProviderTestRequest(
+                provider_type = "custom",
+                base_url = "http://custom.example/v1",
+                model_id = "Qwen/Qwen3-0.6B",
+            ),
+            current_subject = "unsloth",
+        )
+
+    result = _drive(run())
+    assert result.success is True
+    assert result.models_count is None
+    assert captured["init"]["provider_type"] == "custom"
+    assert captured["chat_completion"]["model"] == "Qwen/Qwen3-0.6B"
+    assert captured["chat_completion"]["max_tokens"] == 1
+    assert captured["closed"] is True
+
+
+def test_custom_provider_test_endpoint_requires_model_id(monkeypatch):
+    import importlib.util
+    import sys
+    from pathlib import Path
+
+    module_path = Path(__file__).resolve().parents[1] / "routes" / "providers.py"
+    spec = importlib.util.spec_from_file_location("_providers_route_under_test", module_path)
+    assert spec is not None
+    assert spec.loader is not None
+    providers_route = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = providers_route
+    spec.loader.exec_module(providers_route)
+
+    class _FakeClient:
+        def __init__(self, **kwargs):
+            pass
+
+        async def close(self):
+            pass
+
+    monkeypatch.setattr(providers_route, "ExternalProviderClient", _FakeClient)
+
+    async def run():
+        return await providers_route.test_provider(
+            providers_route.ProviderTestRequest(
+                provider_type = "custom",
+                base_url = "http://custom.example/v1",
+            ),
+            current_subject = "unsloth",
+        )
+
+    result = _drive(run())
+    assert result.success is False
+    assert "model ID" in result.message
+
+
 def test_anthropic_stream_emits_usage_chunk_before_done(monkeypatch):
    sse_events = [
        {
--- a/studio/frontend/src/features/chat/api/providers-api.ts
+++ b/studio/frontend/src/features/chat/api/providers-api.ts
@ -197,6 +197,7 @@ export async function testProviderConnection(payload: {
  providerType: string;
  apiKey: string;
  baseUrl?: string | null;
+  modelId?: string | null;
 }): Promise<ProviderTestResult> {
  return withApiKeyEncryptionRetry(payload.apiKey, async (encryptedApiKey) => {
    const response = await authFetch("/api/providers/test", {
@ -206,6 +207,7 @@ export async function testProviderConnection(payload: {
        provider_type: payload.providerType,
        encrypted_api_key: encryptedApiKey,
        base_url: payload.baseUrl ?? null,
+        model_id: payload.modelId ?? null,
      }),
    });
    return parseJsonOrThrow<ProviderTestResult>(response);
--- a/studio/frontend/src/features/chat/chat-providers-dialog.tsx
+++ b/studio/frontend/src/features/chat/chat-providers-dialog.tsx
@ -59,6 +59,7 @@ import {
  getExternalProviderApiKey,
  isCustomProviderType,
  LEGACY_CUSTOM_PROVIDER_TYPE,
+  CUSTOM_PROVIDER_DISPLAY_NAME,
  removeExternalProviderApiKey,
  setExternalProviderApiKey,
  supportsProviderPromptCaching,
@ -176,7 +177,8 @@ function shouldAppendOpenAiVersionPath(providerType: string): boolean {
  return (
    providerType === "ollama" ||
    providerType === "llama_cpp" ||
-    providerType === "vllm"
+    providerType === "vllm" ||
+    providerType === LEGACY_CUSTOM_PROVIDER_TYPE
  );
 }

@ -230,7 +232,9 @@ export function ChatProvidersSettings({
  const [mutatingProvider, setMutatingProvider] = useState(false);
  const [manualModelIds, setManualModelIds] = useState("");
  const [modelSearchQuery, setModelSearchQuery] = useState("");
-  const [customProviderName, setCustomProviderName] = useState("Custom");
+  const [customProviderName, setCustomProviderName] = useState(
+    CUSTOM_PROVIDER_DISPLAY_NAME,
+  );
  const [isReasoningModel, setIsReasoningModel] = useState(false);
  const reduceMotion = useReducedMotion();
  const connectionsEnabled = useExternalProvidersStore(
@ -939,8 +943,12 @@ export function ChatProvidersSettings({

  async function testProvider(provider: ExternalProviderConfig) {
    const savedKey = getExternalProviderApiKey(provider.id).trim();
-    // Local OpenAI-compat presets skip API keys — run the connection check.
-    if (!savedKey && !supportsRemoteModelCatalog(provider.providerType)) {
+    // Hosted registry providers require keys. Local OpenAI-compatible presets
+    // may be keyless.
+    if (
+      !savedKey &&
+      !supportsRemoteModelCatalog(provider.providerType)
+    ) {
      if (isCustomProviderType(provider.providerType)) {
        await editProvider(provider);
        toast.info(CUSTOM_PROVIDER_MISSING_KEY_MESSAGE);
@ -957,6 +965,10 @@ export function ChatProvidersSettings({
          provider.providerType,
        apiKey: savedKey,
        baseUrl: provider.baseUrl || null,
+        modelId:
+          provider.providerType === LEGACY_CUSTOM_PROVIDER_TYPE
+            ? (provider.models[0] ?? null)
+            : null,
      });
      if (result.success) {
        toast.success(result.message);
@ -1062,6 +1074,16 @@ export function ChatProvidersSettings({
                          </span>
                        </SelectItem>
                      ))}
+                      <SelectItem value={LEGACY_CUSTOM_PROVIDER_TYPE}>
+                        <span className="flex items-center gap-2">
+                          <ApiProviderLogo
+                            providerType={LEGACY_CUSTOM_PROVIDER_TYPE}
+                            className="size-4"
+                            title={CUSTOM_PROVIDER_DISPLAY_NAME}
+                          />
+                          {CUSTOM_PROVIDER_DISPLAY_NAME}
+                        </span>
+                      </SelectItem>
                    </SelectGroup>
                    <SelectSeparator />
                    <SelectGroup>
@ -1144,7 +1166,7 @@ export function ChatProvidersSettings({
                    onChange={(event) =>
                      setCustomProviderName(event.target.value)
                    }
-                    placeholder="Custom"
+                    placeholder={CUSTOM_PROVIDER_DISPLAY_NAME}
                    className="h-9 text-sm"
                  />
                </div>
--- a/studio/frontend/src/features/chat/external-providers.ts
+++ b/studio/frontend/src/features/chat/external-providers.ts
@ -107,6 +107,7 @@ export function providerTypeSupportsVision(

 export const CUSTOM_BACKEND_PROVIDER_TYPE = "openai";
 export const LEGACY_CUSTOM_PROVIDER_TYPE = "custom";
+export const CUSTOM_PROVIDER_DISPLAY_NAME = "Custom";

 export const CUSTOM_PROVIDER_PRESETS = [
  {
@ -130,7 +131,7 @@ export const CUSTOM_PROVIDER_PRESETS = [
 ] as const;

 const CUSTOM_PROVIDER_LABELS: Record<string, string> = {
-  [LEGACY_CUSTOM_PROVIDER_TYPE]: "Custom",
+  [LEGACY_CUSTOM_PROVIDER_TYPE]: CUSTOM_PROVIDER_DISPLAY_NAME,
  ...Object.fromEntries(
    CUSTOM_PROVIDER_PRESETS.map((preset) => [
      preset.providerType,
@ -166,8 +167,9 @@ export function isCustomProviderType(
  return providerType in CUSTOM_PROVIDER_LABELS;
 }

-/** Local OpenAI-compat presets that expose GET /v1/models (no API key). */
+/** OpenAI-compat custom types that may expose GET /v1/models. */
 const REMOTE_MODEL_CATALOG_CUSTOM_PROVIDER_TYPES = new Set([
+  LEGACY_CUSTOM_PROVIDER_TYPE,
  "ollama",
  "vllm",
  "llama_cpp",
@ -189,7 +191,7 @@ export function customPresetSkipsApiKeyField(
  return providerType === "ollama" || providerType === "llama_cpp";
 }

-/** Catalog load plus optional manual model IDs (OpenRouter + local presets). */
+/** Catalog load plus optional manual model IDs. */
 export function allowsManualModelIdsWithCatalog(
  providerType: string | null | undefined,
 ): boolean {
@ -201,7 +203,7 @@ export function allowsManualModelIdsWithCatalog(
 export function customProviderDisplayName(
  providerType: string | null | undefined,
 ): string {
-  if (!providerType) return "Custom";
+  if (!providerType) return CUSTOM_PROVIDER_DISPLAY_NAME;
  return CUSTOM_PROVIDER_LABELS[providerType] ?? providerType;
 }

@ -246,6 +248,11 @@ export function toExternalBackendProviderType(
  if (providerType === "vllm") return "vllm";
  if (providerType === "ollama") return "ollama";
  if (providerType === "llama_cpp") return "llama_cpp";
+  // Generic custom servers are OpenAI-compatible, but should still use the
+  // chat-completions backend path instead of OpenAI's Responses API route.
+  if (providerType === LEGACY_CUSTOM_PROVIDER_TYPE) {
+    return LEGACY_CUSTOM_PROVIDER_TYPE;
+  }
  return isCustomProviderType(providerType)
    ? CUSTOM_BACKEND_PROVIDER_TYPE
    : providerType;