From d8cd95375d189bc558d1e3d525bf71d9892a6d31 Mon Sep 17 00:00:00 2001 From: Ned Martorell <32619850+nmartorell@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:39:34 -0400 Subject: [PATCH 01/68] feat(config): require explicit Dataiku instance types --- .dataiku/http-config.json.entra-example | 1 + .../http-config.json.generic_oidc-example | 1 + .dataiku/stdio-config.json.example | 2 + .env.example | 1 + CODING_STANDARDS_AND_STRUCTURE.md | 1 + README.md | 18 +- dataiku_mcp/config/models.py | 7 + dataiku_mcp/config/request.py | 1 + dataiku_mcp/config/stdio.py | 24 ++- dataiku_mcp/setup_server.py | 23 ++- dataiku_mcp/tools/instances.py | 5 +- docs/capabilities.md | 4 +- docs/http-deployment.md | 11 +- skills/dataiku-headless-setup/SKILL.md | 4 +- skills/dataiku-headless/SKILL.md | 2 +- tests/test_config.py | 156 +++++++++++++++++- tests/test_http_config.py | 62 +++++++ tests/test_instance_pinning.py | 16 +- tests/test_instances.py | 46 +++++- tests/test_setup_server.py | 68 ++++++++ 20 files changed, 428 insertions(+), 25 deletions(-) diff --git a/.dataiku/http-config.json.entra-example b/.dataiku/http-config.json.entra-example index 451f482e..b80e54e0 100644 --- a/.dataiku/http-config.json.entra-example +++ b/.dataiku/http-config.json.entra-example @@ -15,6 +15,7 @@ }, "dss_instances": { "prod": { + "instance_type": "automation", "url": "https://dataiku.example", "delegated_scope": "api://replace-with-dataiku-app-client-id/dataiku.access" } diff --git a/.dataiku/http-config.json.generic_oidc-example b/.dataiku/http-config.json.generic_oidc-example index 34fb7a7b..dfe6304f 100644 --- a/.dataiku/http-config.json.generic_oidc-example +++ b/.dataiku/http-config.json.generic_oidc-example @@ -23,6 +23,7 @@ }, "dss_instances": { "prod": { + "instance_type": "automation", "url": "https://dataiku.example", "delegated_audience": "dataiku-prod", "delegated_scope": "dataiku.api" diff --git a/.dataiku/stdio-config.json.example b/.dataiku/stdio-config.json.example index 79776585..bc9fe8cb 100644 --- a/.dataiku/stdio-config.json.example +++ b/.dataiku/stdio-config.json.example @@ -2,11 +2,13 @@ "default_instance": "dev", "dss_instances": { "dev": { + "instance_type": "design", "url": "https://dev.dataiku.com", "api_key": "your-dev-api-key", "no_check_certificate": false }, "prod": { + "instance_type": "automation", "url": "https://prod.dataiku.com", "api_key": "your-prod-api-key", "no_check_certificate": false diff --git a/.env.example b/.env.example index 8140dc99..1bcfd520 100644 --- a/.env.example +++ b/.env.example @@ -4,6 +4,7 @@ # the configure_instance tool instead. When set, this instance takes precedence # over the resolved config file's `default_instance`. DKU_INSTANCE_NAME="dev-dataiku" +DKU_INSTANCE_TYPE=design # design, automation, deployer, or agent-management DKU_DSS_URL=https://your-instance.dataiku.com DKU_NO_CHECK_CERTIFICATE=false # Set to true/1 to skip SSL verification DKU_API_KEY=your-api-key diff --git a/CODING_STANDARDS_AND_STRUCTURE.md b/CODING_STANDARDS_AND_STRUCTURE.md index 60506919..01d7e323 100644 --- a/CODING_STANDARDS_AND_STRUCTURE.md +++ b/CODING_STANDARDS_AND_STRUCTURE.md @@ -20,6 +20,7 @@ Set environment variables: ```bash export DKU_DSS_URL="https://your-instance.dataiku.com" +export DKU_INSTANCE_TYPE="design" export DKU_API_KEY="your-api-key" ``` diff --git a/README.md b/README.md index 3da40140..e1cf3c04 100644 --- a/README.md +++ b/README.md @@ -166,7 +166,7 @@ The onboarding flow is: 1. Ask the agent to **Set up Dataiku Headless** (or run `/dataiku-headless:dataiku-headless-setup` in Claude Code). 2. Approve the MCP URL prompt. -3. Enter an instance name, Dataiku URL, and personal API key. +3. Enter an instance name, choose its type, and supply its Dataiku URL and personal API key. 4. Repeat to add more instances; use `list_instances` and `switch_instance` while working. The API key never appears in MCP tool arguments. @@ -188,12 +188,26 @@ The server loads environment and profile settings at startup. Profile additions deletions refresh both the resolved file and the in-memory catalog; otherwise, manual or environment changes require a restart. See [`.dataiku/stdio-config.json.example`](.dataiku/stdio-config.json.example) for the file shape. +Profiles require an explicit `instance_type`: `design`, `automation`, `deployer`, or +`agent-management`. Choose the type on the setup page or set it in a JSON profile. +Environment credentials require `DKU_INSTANCE_TYPE` whenever `DKU_DSS_URL` is set. +Missing, empty, or unsupported values are rejected for setup submissions, +environment settings, and configuration profiles. An environment instance keeps +its own type and takes precedence over a profile with the same name. + +Instance tools report this configured type. It is operator-supplied metadata for +future routing; it is not discovered or verified through the API and does not change +client selection, connection testing, or tool availability. Choosing Agent Management +does not add client support. Configuration changes require a restart unless saved +through the setup page. + Environment variables are an explicit override: **.env file:** Copy `.env.example` to `.env` and fill in your values: ```bash DKU_DSS_URL=https://your-instance.dataiku.com +DKU_INSTANCE_TYPE=design DKU_API_KEY=your-api-key DKU_MCP_MAX_WORKERS=4 DKU_NO_CHECK_CERTIFICATE=false @@ -210,7 +224,7 @@ Put instance info in the resolved configuration file. See `.dataiku/stdio-config After adding multiple instance configs, you can use the `list_instances`, `switch_instance`, and `get_current_instance` MCP tools to manage instances from the agent. Auth resolution order: -1. Environment variables: `DKU_DSS_URL`, `DKU_API_KEY`, and optional `DKU_NO_CHECK_CERTIFICATE` +1. Environment variables: `DKU_DSS_URL`, `DKU_API_KEY`, `DKU_INSTANCE_TYPE`, and optional `DKU_NO_CHECK_CERTIFICATE` 2. The resolved configuration file, using its `default_instance` ## Run diff --git a/dataiku_mcp/config/models.py b/dataiku_mcp/config/models.py index 8e5d2118..dc797986 100644 --- a/dataiku_mcp/config/models.py +++ b/dataiku_mcp/config/models.py @@ -26,6 +26,9 @@ ) +InstanceType = Literal["design", "automation", "deployer", "agent-management"] + + @dataclass(frozen=True) class DSSInstance: name: str @@ -33,6 +36,7 @@ class DSSInstance: api_key: str = field(repr=False) no_check_certificate: bool source: str + instance_type: InstanceType description: str = "" delegated_audience: str = "" delegated_scope: str = "" @@ -53,6 +57,7 @@ class _DSSInstanceConfig(_StrictConfigModel): url: NonEmptyString no_check_certificate: bool = False description: str = "" + instance_type: InstanceType class StdioDSSInstanceConfig(_DSSInstanceConfig): @@ -71,6 +76,7 @@ def to_instance( no_check_certificate=self.no_check_certificate, source=source, description=self.description, + instance_type=self.instance_type, ) @@ -168,6 +174,7 @@ def to_instance(self, name: str) -> DSSInstance: description=self.description, delegated_audience=self.delegated_audience or "", delegated_scope=self.delegated_scope, + instance_type=self.instance_type, ) diff --git a/dataiku_mcp/config/request.py b/dataiku_mcp/config/request.py index ec69a6a3..a8593748 100644 --- a/dataiku_mcp/config/request.py +++ b/dataiku_mcp/config/request.py @@ -140,4 +140,5 @@ def set_current_instance(name: str) -> dict: "name": selected.name, "url": selected.url, "description": selected.description, + "instance_type": selected.instance_type, } diff --git a/dataiku_mcp/config/stdio.py b/dataiku_mcp/config/stdio.py index c0682842..06094cda 100644 --- a/dataiku_mcp/config/stdio.py +++ b/dataiku_mcp/config/stdio.py @@ -22,7 +22,7 @@ from pydantic import ValidationError from .files import read_json_object, write_json_atomic -from .models import DSSInstance, StdioConfig, StdioDSSInstanceConfig +from .models import DSSInstance, InstanceType, StdioConfig, StdioDSSInstanceConfig DEFAULT_SETTINGS_PATH = Path.home() / ".dataiku" / "stdio-config.json" @@ -111,14 +111,21 @@ def _load_instance_from_env_vars() -> DSSInstance | None: if not os.environ.get("DKU_DSS_URL"): return None + if "DKU_INSTANCE_TYPE" not in os.environ: + raise ValueError("DKU_INSTANCE_TYPE is required when DKU_DSS_URL is set.") + no_check_certificate = os.environ.get("DKU_NO_CHECK_CERTIFICATE", "").strip() try: - instance = StdioDSSInstanceConfig( - url=os.environ["DKU_DSS_URL"], - api_key=os.environ.get("DKU_API_KEY", ""), - no_check_certificate=( - bool(no_check_certificate) and no_check_certificate.lower() != "false" - ), + instance = StdioDSSInstanceConfig.model_validate( + { + "url": os.environ["DKU_DSS_URL"], + "api_key": os.environ.get("DKU_API_KEY", ""), + "instance_type": os.environ["DKU_INSTANCE_TYPE"], + "no_check_certificate": ( + bool(no_check_certificate) + and no_check_certificate.lower() != "false" + ), + } ) except ValidationError as err: raise ValueError(f"Invalid stdio environment settings: {err}") from None @@ -194,6 +201,7 @@ def add_instance_to_config( api_key: str, *, description: str = "", + instance_type: InstanceType, no_check_certificate: bool = False, set_default: bool = False, ) -> dict: @@ -204,6 +212,7 @@ def add_instance_to_config( url=url, api_key=api_key, description=description, + instance_type=instance_type, no_check_certificate=no_check_certificate, ) config = _load_config() @@ -216,6 +225,7 @@ def add_instance_to_config( "name": name, "url": url, "description": description, + "instance_type": instance.instance_type, "path": str(get_settings_path()), "default_instance": config.default_instance, } diff --git a/dataiku_mcp/setup_server.py b/dataiku_mcp/setup_server.py index 59f5f815..58c55590 100644 --- a/dataiku_mcp/setup_server.py +++ b/dataiku_mcp/setup_server.py @@ -20,11 +20,13 @@ from dataclasses import dataclass from html import escape from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import get_args from urllib.parse import parse_qs, urlsplit import dataikuapi from .config import request, stdio +from .config.models import InstanceType SESSION_LIFETIME_SECONDS = 10 * 60 MAX_REQUEST_BYTES = 16 * 1024 @@ -62,6 +64,12 @@ def _validate_form(form: dict[str, list[str]]) -> dict: url = form.get("url", [""])[0].strip() api_key = form.get("api_key", [""])[0].strip() description = form.get("description", [""])[0].strip() + instance_type = form.get("instance_type", [""])[0] + + if instance_type not in get_args(InstanceType): + raise ValueError( + "Instance type must be one of: " + ", ".join(get_args(InstanceType)) + "." + ) if not name: raise ValueError("Instance name is required.") @@ -93,6 +101,7 @@ def _validate_form(form: dict[str, list[str]]) -> dict: "url": url, "api_key": api_key, "description": description, + "instance_type": instance_type, "no_check_certificate": "no_check_certificate" in form, "set_default": "set_default" in form, } @@ -115,6 +124,13 @@ def _page( values = values or {} name = escape(values.get("name", "")) description = escape(values.get("description", "")) + instance_type = values.get("instance_type", "") + instance_type_options = "".join( + f'" + for value in get_args(InstanceType) + ) url = escape(values.get("url", "")) api_key = escape(values.get("api_key", "")) set_default = " checked" if is_initial_page or values.get("set_default") else "" @@ -146,9 +162,9 @@ def _page( .grid {{ display:grid; grid-template-columns:1fr 1fr; gap:18px; }} label {{ display:block; font-size:13px; font-weight:650; margin-bottom:8px; }} .full {{ grid-column:1 / -1; }} - input[type=text], input[type=url], input[type=password] {{ width:100%; border:1px solid var(--line); border-radius:8px; padding:11px 12px; color:var(--ink); background:var(--paper); font:inherit; outline:none; transition:border-color 150ms ease, box-shadow 150ms ease, background-color 150ms ease; }} - input:hover {{ border-color:#aeb9b9; }} - input:focus-visible {{ border-color:var(--teal-dark); box-shadow:0 0 0 3px #00a6a626; background:var(--surface); }} + input[type=text], input[type=url], input[type=password], select {{ width:100%; border:1px solid var(--line); border-radius:8px; padding:11px 12px; color:var(--ink); background:var(--paper); font:inherit; outline:none; transition:border-color 150ms ease, box-shadow 150ms ease, background-color 150ms ease; }} + input:hover, select:hover {{ border-color:#aeb9b9; }} + input:focus-visible, select:focus-visible {{ border-color:var(--teal-dark); box-shadow:0 0 0 3px #00a6a626; background:var(--surface); }} .hint {{ color:var(--muted); font-size:12px; margin-top:4px; }} .checks {{ display:grid; gap:12px; margin:20px 0 24px; }} .check {{ display:flex; gap:10px; align-items:flex-start; font-weight:500; margin:0; }} @@ -181,6 +197,7 @@ def _page(
A short name used when switching instances.
+
Enter the URL of your Dataiku instance.
Create one in Dataiku under Profile & Settings → API keys.
diff --git a/dataiku_mcp/tools/instances.py b/dataiku_mcp/tools/instances.py index 0ed0a666..9b20cd3e 100644 --- a/dataiku_mcp/tools/instances.py +++ b/dataiku_mcp/tools/instances.py @@ -60,9 +60,12 @@ async def list_instances(ctx: Context) -> str: "url": inst.url, "description": inst.description, "active": name == current_instance_name, + "instance_type": inst.instance_type, } ) - return compact_json(columnar(result, ["name", "url", "description", "active"])) + return compact_json( + columnar(result, ["name", "url", "description", "active", "instance_type"]) + ) @mcp.tool( diff --git a/docs/capabilities.md b/docs/capabilities.md index 3c389167..1e79feb5 100644 --- a/docs/capabilities.md +++ b/docs/capabilities.md @@ -105,7 +105,7 @@ Local client configuration, not Dataiku objects. | Tool | Effect | |---|---| -| `configure_instance` | Connect a local stdio instance; opens a local page for URL + API key (disabled in HTTP mode) | +| `configure_instance` | Connect a local stdio instance; opens a local page for URL, instance type, and API key (disabled in HTTP mode) | | `switch_instance` | Change the active instance; in HTTP mode saves only the authenticated user’s selected catalog instance | -| `list_instances`, `get_current_instance` | Show configured instances and the active one; `get_current_instance` includes the Dataiku version when available | +| `list_instances`, `get_current_instance` | Show configured instances, their declared types, and the active one; `get_current_instance` includes the Dataiku version when available | | `delete_instance` | Removes a **saved connection profile from the local config file**. Does not touch the Dataiku instance. | diff --git a/docs/http-deployment.md b/docs/http-deployment.md index 4739d5b2..c6cbf4df 100644 --- a/docs/http-deployment.md +++ b/docs/http-deployment.md @@ -105,10 +105,17 @@ The settings file has four sections: - `auth` configures incoming-token verification, optional interactive login, and token exchange. - `dss_instances` lists the Dataiku endpoints users may select and the audience or - scope requested for each one. + scope requested for each one. Each entry requires an explicit `instance_type`: + `design`, `automation`, `deployer`, or `agent-management`. - `user_selections` records each authenticated user's current instance. Start with an empty object; the server manages it. +Instance types are operator-supplied metadata reported by instance tools. They do +not trigger API discovery or change client selection, connection testing, or tool +availability. Agent Management client routing is not implemented by this field. +Existing HTTP catalogs must be updated by the administrator; missing types fail +validation without automatically migrating or rewriting the file. + Authentication, transport, and instance settings are loaded at startup, so restart the server after changing them. User selections are updated in memory and persisted as users call `switch_instance`. This file-backed selection state supports one @@ -228,6 +235,7 @@ and the [OBO flow](https://learn.microsoft.com/en-us/entra/identity-platform/v2- "dss_instances": { "prod": { "url": "https://dataiku.example", + "instance_type": "automation", "delegated_scope": "api://replace-with-dataiku-app-client-id/dataiku.access" } }, @@ -293,6 +301,7 @@ setup, trust, and consent requirements depend on the authorization server. "dss_instances": { "prod": { "url": "https://dataiku.example", + "instance_type": "automation", "delegated_audience": "dataiku-prod", "delegated_scope": "dataiku.api" } diff --git a/skills/dataiku-headless-setup/SKILL.md b/skills/dataiku-headless-setup/SKILL.md index 8ab3041c..3561405c 100644 --- a/skills/dataiku-headless-setup/SKILL.md +++ b/skills/dataiku-headless-setup/SKILL.md @@ -22,7 +22,7 @@ Bring a new or broken plugin installation to a verified Dataiku connection. A re A startup line on stderr followed by exit status 0 is expected. Do not substitute `uv sync --locked --script`: it caches downloads but leaves environment creation for the first server launch. 7. Retry `list_instances`. If the tools are still unavailable, reload the plugin or restart the agent once, then retry before diagnosing a Dataiku connection problem. -8. When `list_instances` succeeds, if no instance is configured, run `configure_instance` and have the user complete the local setup page. If multiple instances exist without an active one, ask which to use and run `switch_instance`. +8. When `list_instances` succeeds, if no instance is configured, run `configure_instance` and have the user complete the local setup page, explicitly choosing `design`, `automation`, `deployer`, or `agent-management`. Environment credentials require `DKU_INSTANCE_TYPE`; all configuration profiles require `instance_type`. This is declared metadata, not API discovery or client routing, and choosing a type does not add corresponding client functionality. If multiple instances exist without an active one, ask which to use and run `switch_instance`. 9. Verify the active profile with `get_current_instance`. If its `connection_status` is `failed`, the discovered profile cannot connect to Dataiku; run `configure_instance` to replace or add a working profile instead of treating it as set up. When it is `connected`, make a lightweight read-only Dataiku call such as `list_projects` to validate the available Dataiku access. Never request or repeat the API key in chat. -Report completion with the active instance name and URL, Dataiku version when available, and whether the Dataiku read succeeded. Include the uv version only when runtime recovery was needed. If a restart is required, say that setup is incomplete and give the single next action. +Report completion with the active instance name, URL, and configured type, Dataiku version when available, and whether the Dataiku read succeeded. Include the uv version only when runtime recovery was needed. If a restart is required, say that setup is incomplete and give the single next action. diff --git a/skills/dataiku-headless/SKILL.md b/skills/dataiku-headless/SKILL.md index c2340228..93367986 100644 --- a/skills/dataiku-headless/SKILL.md +++ b/skills/dataiku-headless/SKILL.md @@ -10,7 +10,7 @@ Use this for any Dataiku task. Choose the right reference guide first, inspect t ## Shared Operating Rules 1. If the user asks to install, set up, connect, or repair Dataiku Headless, or its MCP tools are unavailable just after installation, read `../dataiku-headless-setup/SKILL.md`. Choose either local stdio or customer-managed HTTP setup and never enable both. -2. Ensure an instance is configured before any Dataiku work. In local stdio mode, if `get_current_instance` errors or `list_instances` is empty, run `configure_instance` first. In HTTP mode, use `list_instances` then `switch_instance`; the instance catalog is platform-managed. Keep the reported `dataiku_version` in context for version-sensitive requests. +2. Ensure an instance is configured before any Dataiku work. In local stdio mode, if `get_current_instance` errors or `list_instances` is empty, run `configure_instance` first. In HTTP mode, use `list_instances` then `switch_instance`; the instance catalog is platform-managed. Keep the reported `dataiku_version` in context for version-sensitive requests. The required `instance_type` is configured metadata, not API verification or a guarantee of client support. 3. Discover project keys and object identifiers through tools; do not invent them. 4. Read before write. Inspect the current object, flow context, jobs, or run history before changing anything. 5. Treat the matching reference guide as the source of truth for object-specific concepts, inspection steps, and required references. diff --git a/tests/test_config.py b/tests/test_config.py index 19cd5059..bfbf8d77 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -20,7 +20,12 @@ import pytest from dataiku_mcp.config import request, stdio -from dataiku_mcp.config.models import DSSInstance, StdioConfig, StdioDSSInstanceConfig +from dataiku_mcp.config.models import ( + DSSInstance, + HTTPDSSInstanceConfig, + StdioConfig, + StdioDSSInstanceConfig, +) from dataiku_mcp.tools import instances as instance_tools @@ -34,6 +39,7 @@ def isolated_config(tmp_path, monkeypatch): "DKU_DSS_URL", "DKU_API_KEY", "DKU_INSTANCE_NAME", + "DKU_INSTANCE_TYPE", "DKU_NO_CHECK_CERTIFICATE", "DKU_CONFIG_FILE", ): @@ -45,6 +51,7 @@ def add_instance(name: str, *, set_default: bool = False) -> None: name, f"https://{name}.example.com", f"{name}-api-key", + instance_type="design", set_default=set_default, ) @@ -89,7 +96,11 @@ def test_stdio_config_migrates_legacy_cwd_settings(tmp_path, monkeypatch, caplog { "default_instance": "dev", "dss_instances": { - "dev": {"url": "https://dev.example.com", "api_key": "api-key"} + "dev": { + "url": "https://dev.example.com", + "api_key": "api-key", + "instance_type": "automation", + } }, } ) @@ -105,6 +116,10 @@ def test_stdio_config_migrates_legacy_cwd_settings(tmp_path, monkeypatch, caplog assert not legacy_path.exists() assert stdio._load_config().default_instance == "dev" + assert ( + json.loads(canonical_path.read_text())["dss_instances"]["dev"]["instance_type"] + == "automation" + ) assert "Migrated legacy stdio configuration" in caplog.text assert "api-key" not in caplog.text @@ -205,6 +220,52 @@ def test_stdio_getters_require_initialization(): stdio.get_instances() +def test_stdio_config_requires_explicit_type_without_rewriting_file(): + path = stdio.get_settings_path() + path.write_text( + json.dumps( + { + "dss_instances": { + "dev": {"url": "https://dev.example.com", "api_key": "secret-key"} + } + } + ) + ) + original = path.read_bytes() + + with pytest.raises(ValueError, match="instance_type") as error: + stdio.initialize_config() + assert path.read_bytes() == original + assert "secret-key" not in str(error.value) + + +@pytest.mark.parametrize( + "config_model", [StdioDSSInstanceConfig, HTTPDSSInstanceConfig] +) +def test_instance_models_require_explicit_types(config_model): + fields = {"url": "https://typed.example"} + if config_model is StdioDSSInstanceConfig: + fields["api_key"] = "secret" + else: + fields["delegated_scope"] = "dss.api" + with pytest.raises(ValueError, match="instance_type"): + config_model.model_validate(fields) + + +def test_runtime_instance_requires_explicit_type(): + with pytest.raises(TypeError, match="instance_type"): + DSSInstance("typed", "https://typed.example", "secret", False, "config") + + +def test_environment_requires_explicit_type_without_inheriting_profile(monkeypatch): + add_instance("typed", set_default=True) + monkeypatch.setenv("DKU_INSTANCE_NAME", "typed") + monkeypatch.setenv("DKU_DSS_URL", "https://typed.example") + monkeypatch.setenv("DKU_API_KEY", "secret-key") + with pytest.raises(ValueError, match="DKU_INSTANCE_TYPE is required"): + stdio.initialize_config() + + def test_setting_stdio_path_invalidates_cached_state(): add_instance("dev", set_default=True) stdio.initialize_config() @@ -256,6 +317,8 @@ def test_stdio_config_rejects_unknown_instance_fields(unknown_field): "dss_instances": { "dev": { "url": "https://dev.example.com", + "api_key": "api-key", + "instance_type": "design", unknown_field: "unexpected", } } @@ -283,7 +346,7 @@ def test_stdio_config_rejects_unknown_instance_fields(unknown_field): ) def test_stdio_config_rejects_invalid_instance_fields(instance): stdio.get_settings_path().write_text( - json.dumps({"dss_instances": {"dev": instance}}) + json.dumps({"dss_instances": {"dev": {"instance_type": "design", **instance}}}) ) with pytest.raises(ValueError, match="Invalid stdio settings"): @@ -308,6 +371,7 @@ def test_stdio_config_validation_errors_do_not_expose_api_keys(): "dev": { "url": "https://dev.example.com", "api_key": [secret], + "instance_type": "design", } } } @@ -330,6 +394,7 @@ def test_stdio_config_example_is_valid(monkeypatch): assert config.dss_instances["dev"] == StdioDSSInstanceConfig( url="https://dev.dataiku.com", api_key="your-dev-api-key", + instance_type="design", no_check_certificate=False, ) @@ -338,6 +403,7 @@ def test_stdio_instance_config_converts_to_runtime_instance(): config = StdioDSSInstanceConfig( url="https://dev.example.com", api_key="api-key", + instance_type="design", no_check_certificate=True, description="Development", ) @@ -348,6 +414,7 @@ def test_stdio_instance_config_converts_to_runtime_instance(): api_key="api-key", no_check_certificate=True, source="config", + instance_type="design", description="Development", ) @@ -359,14 +426,91 @@ def test_runtime_instance_repr_hides_api_key(): api_key="do-not-print-this-api-key", no_check_certificate=False, source="config", + instance_type="design", ) assert "do-not-print-this-api-key" not in repr(instance) + assert instance.instance_type == "design" + + +@pytest.mark.parametrize( + "instance_type", ["design", "automation", "deployer", "agent-management"] +) +def test_stdio_instance_type_survives_save_load_and_selection( + monkeypatch, instance_type +): + def unexpected_client(*args, **kwargs): + pytest.fail("Instance configuration must not construct a DSS client") + + monkeypatch.setattr("dataikuapi.DSSClient", unexpected_client) + result = stdio.add_instance_to_config( + "typed", "https://typed.example", "secret", instance_type=instance_type + ) + stdio.initialize_config() + selected = request.set_current_instance("typed") + saved = json.loads(stdio.get_settings_path().read_text()) + + assert result["instance_type"] == instance_type + assert selected["instance_type"] == instance_type + assert stdio.get_instances()["typed"].instance_type == instance_type + assert stdio.get_current_instance().instance_type == instance_type + assert saved["dss_instances"]["typed"]["instance_type"] == instance_type + + +@pytest.mark.parametrize("instance_type", ["", "govern", "DESIGN", " design", None, 1]) +def test_stdio_config_rejects_invalid_instance_types(instance_type): + stdio.get_settings_path().write_text( + json.dumps( + { + "dss_instances": { + "typed": { + "url": "https://typed.example", + "api_key": "secret-key", + "instance_type": instance_type, + } + } + } + ) + ) + with pytest.raises(ValueError, match="instance_type") as error: + stdio.initialize_config() + assert "secret-key" not in str(error.value) + + +@pytest.mark.parametrize( + "instance_type", ["design", "automation", "deployer", "agent-management"] +) +def test_environment_instance_type_overrides_profile(monkeypatch, instance_type): + stdio.add_instance_to_config( + "typed", "https://profile.example", "profile-key", instance_type="deployer" + ) + monkeypatch.setenv("DKU_INSTANCE_NAME", "typed") + monkeypatch.setenv("DKU_DSS_URL", "https://environment.example") + monkeypatch.setenv("DKU_API_KEY", "environment-key") + monkeypatch.setenv("DKU_INSTANCE_TYPE", instance_type) + stdio.initialize_config() + + instance = stdio.get_instances()["typed"] + assert instance.source == "environment" + assert instance.instance_type == instance_type + assert stdio.get_current_instance() == instance + + +@pytest.mark.parametrize("instance_type", ["", "govern", "DESIGN", " design"]) +def test_environment_rejects_invalid_instance_type(monkeypatch, instance_type): + monkeypatch.setenv("DKU_DSS_URL", "https://typed.example") + monkeypatch.setenv("DKU_API_KEY", "secret-key") + monkeypatch.setenv("DKU_INSTANCE_TYPE", instance_type) + + with pytest.raises(ValueError, match="Invalid stdio environment settings") as error: + stdio.initialize_config() + assert "secret-key" not in str(error.value) @pytest.mark.parametrize("api_key", [None, ""]) def test_stdio_environment_requires_api_key(monkeypatch, api_key): monkeypatch.setenv("DKU_DSS_URL", "https://dev.example.com") + monkeypatch.setenv("DKU_INSTANCE_TYPE", "design") if api_key is not None: monkeypatch.setenv("DKU_API_KEY", api_key) @@ -376,6 +520,7 @@ def test_stdio_environment_requires_api_key(monkeypatch, api_key): def test_stdio_environment_uses_validated_instance(monkeypatch): monkeypatch.setenv("DKU_DSS_URL", "https://dev.example.com") + monkeypatch.setenv("DKU_INSTANCE_TYPE", "design") monkeypatch.setenv("DKU_API_KEY", "api-key") stdio.initialize_config() @@ -390,6 +535,7 @@ def test_stdio_getters_use_the_startup_snapshot(monkeypatch): add_instance("dev", set_default=True) stdio.initialize_config() monkeypatch.setenv("DKU_DSS_URL", "https://environment.example.com") + monkeypatch.setenv("DKU_INSTANCE_TYPE", "design") monkeypatch.setenv("DKU_API_KEY", "environment-key") stdio.get_settings_path().write_text( json.dumps( @@ -399,6 +545,7 @@ def test_stdio_getters_use_the_startup_snapshot(monkeypatch): "changed": { "url": "https://changed.example.com", "api_key": "changed-key", + "instance_type": "design", } }, } @@ -427,6 +574,7 @@ def fail_save(_config): "prod", "https://prod.example.com", "prod-key", + instance_type="design", ) assert set(stdio.get_instances()) == {"dev"} @@ -466,6 +614,7 @@ def test_stdio_config_save_excludes_runtime_fields(): "dev", "https://dev.example.com", "api-key", + instance_type="design", description="Development", no_check_certificate=True, set_default=True, @@ -479,6 +628,7 @@ def test_stdio_config_save_excludes_runtime_fields(): "no_check_certificate": True, "description": "Development", "api_key": "api-key", + "instance_type": "design", } assert "name" not in instance assert "source" not in instance diff --git a/tests/test_http_config.py b/tests/test_http_config.py index 57987d73..2d40155a 100644 --- a/tests/test_http_config.py +++ b/tests/test_http_config.py @@ -70,11 +70,13 @@ def http_config(monkeypatch, tmp_path): }, "dss_instances": { "sandbox": { + "instance_type": "design", "url": "https://sandbox.example", "delegated_audience": "dss-sandbox", "delegated_scope": "dss.api", }, "prod": { + "instance_type": "automation", "url": "https://prod.example", "delegated_audience": "dss-prod", "delegated_scope": "dss.api", @@ -174,6 +176,7 @@ def test_http_config_rejects_api_keys(http_config): def test_http_instance_config_converts_to_runtime_instance(): config = HTTPDSSInstanceConfig( + instance_type="automation", url="https://prod.example.com", delegated_audience="dss-prod", delegated_scope="dss.api", @@ -187,12 +190,67 @@ def test_http_instance_config_converts_to_runtime_instance(): api_key="", no_check_certificate=True, source="http", + instance_type="automation", description="Production", delegated_audience="dss-prod", delegated_scope="dss.api", ) +@pytest.mark.parametrize( + "instance_type", ["design", "automation", "deployer", "agent-management"] +) +def test_http_instance_type_survives_selection_save( + http_config, monkeypatch, instance_type +): + document = json.loads(http_config.read_text()) + document["dss_instances"]["prod"]["instance_type"] = instance_type + http_config.write_text(json.dumps(document)) + + def unexpected_client(*args, **kwargs): + pytest.fail("Instance configuration must not construct a DSS client") + + monkeypatch.setattr("dataikuapi.DSSClient", unexpected_client) + http.initialize_config() + identity = request.bind_http_identity("https://idp.example", "alice") + try: + result = request.set_current_instance("prod") + pinned = request.pin_current_instance() + try: + assert request.get_pinned_instance().instance_type == instance_type + finally: + request.reset_pinned_instance(pinned) + finally: + request.reset_http_identity(identity) + + saved = json.loads(http_config.read_text()) + assert result["instance_type"] == instance_type + assert saved["dss_instances"]["prod"]["instance_type"] == instance_type + assert http._load_config().dss_instances["prod"].instance_type == instance_type + + +@pytest.mark.parametrize("instance_type", ["", "govern", "DESIGN", " design", None, 1]) +def test_http_config_rejects_invalid_instance_types(http_config, instance_type): + document = json.loads(http_config.read_text()) + document["dss_instances"]["prod"]["instance_type"] = instance_type + http_config.write_text(json.dumps(document)) + + with pytest.raises(ValueError, match="instance_type") as error: + http.initialize_config() + assert "exchange-secret" not in str(error.value) + + +def test_http_requires_explicit_type_without_migrating_catalog(http_config): + document = json.loads(http_config.read_text()) + document["dss_instances"]["prod"].pop("instance_type") + http_config.write_text(json.dumps(document)) + original = http_config.read_bytes() + + with pytest.raises(ValueError, match="instance_type"): + http.initialize_config() + assert http_config.read_bytes() == original + + def test_http_config_rejects_obsolete_user_defaults_key(http_config): document = json.loads(http_config.read_text()) document["user_defaults"] = document.pop("user_selections") @@ -447,6 +505,7 @@ def test_generic_oidc_http_config_example_is_valid(monkeypatch): assert isinstance(http._load_config(), HTTPConfig) assert set(instances) == {"prod"} assert instances["prod"].url == "https://dataiku.example" + assert instances["prod"].instance_type == "automation" assert instances["prod"].delegated_audience == "dataiku-prod" assert instances["prod"].delegated_scope == "dataiku.api" assert selections == {} @@ -473,6 +532,7 @@ def test_entra_http_config_example_is_valid(monkeypatch): ) instances, selections = http.get_instances_and_selections() assert instances["prod"].delegated_audience == "" + assert instances["prod"].instance_type == "automation" assert instances["prod"].delegated_scope == ( "api://replace-with-dataiku-app-client-id/dataiku.access" ) @@ -765,6 +825,7 @@ def json(self): False, "http", delegated_audience="dss-prod", + instance_type="automation", delegated_scope="dss.api", ), ) @@ -850,6 +911,7 @@ def json(self): False, "http", delegated_audience="dss-prod", + instance_type="automation", delegated_scope="dss.api", ), ) diff --git a/tests/test_instance_pinning.py b/tests/test_instance_pinning.py index fdc78531..25c1308a 100644 --- a/tests/test_instance_pinning.py +++ b/tests/test_instance_pinning.py @@ -39,10 +39,20 @@ def get_project(self, project_key: str) -> _FakeProject: def test_tool_keeps_its_initial_instance_after_a_concurrent_switch(monkeypatch): instance_a = DSSInstance( - "instance-a", "https://a.example", "key-a", False, "config" + "instance-a", + "https://a.example", + "key-a", + False, + "config", + instance_type="automation", ) instance_b = DSSInstance( - "instance-b", "https://b.example", "key-b", False, "config" + "instance-b", + "https://b.example", + "key-b", + False, + "config", + instance_type="deployer", ) observed_instances = [] fake_client = _FakeClient() @@ -52,6 +62,7 @@ def test_tool_keeps_its_initial_instance_after_a_concurrent_switch(monkeypatch): def get_client(): stdio._current_instance = instance_b observed_instances.append(request.get_pinned_instance().name) + assert request.get_pinned_instance().instance_type == "automation" return fake_client monkeypatch.setattr(projects, "get_dss_client", get_client) @@ -68,3 +79,4 @@ async def call_tool(): assert not result.is_error assert observed_instances == ["instance-a"] assert stdio.get_current_instance() == instance_b + assert stdio.get_current_instance().instance_type == "deployer" diff --git a/tests/test_instances.py b/tests/test_instances.py index 0a891775..9bb6eba3 100644 --- a/tests/test_instances.py +++ b/tests/test_instances.py @@ -16,6 +16,9 @@ import asyncio import json +from dataclasses import replace + +import pytest from dataiku_mcp.config import request from dataiku_mcp.config.models import DSSInstance @@ -43,7 +46,13 @@ def get_instance_info(self) -> _FakeInstanceInfo: def _instance(name: str) -> DSSInstance: return DSSInstance( - name, f"https://{name}.example", API_KEY, False, "config", f"{name} desc" + name, + f"https://{name}.example", + API_KEY, + False, + "config", + "design", + f"{name} desc", ) @@ -58,6 +67,40 @@ def _current(monkeypatch, client: _FakeClient) -> dict: return json.loads(asyncio.run(instances.get_current_instance(FakeContext()))) +@pytest.mark.parametrize( + "instance_type", ["design", "automation", "deployer", "agent-management"] +) +def test_instance_tools_report_configured_type(monkeypatch, instance_type): + configured = replace(_instance("primary"), instance_type=instance_type) + monkeypatch.setattr(request, "get_instances", lambda: {"primary": configured}) + monkeypatch.setattr(request, "get_pinned_instance", lambda: configured) + + def unexpected_client(): + pytest.fail("Listing and switching must not construct a DSS client") + + monkeypatch.setattr(instances, "get_dss_client", unexpected_client) + monkeypatch.setattr( + "dataiku_mcp.config.stdio.set_current_instance", lambda inst: None + ) + monkeypatch.setattr(request, "is_http_request", lambda: False) + listing = asyncio.run(instances.list_instances(FakeContext())) + switching = asyncio.run(instances.switch_instance("primary", FakeContext())) + table = json.loads(listing) + assert table["columns"] == ["name", "url", "description", "active", "instance_type"] + assert table["rows"][0][-1] == instance_type + assert json.loads(switching)["instance_type"] == instance_type + + # API information is observational; it must not replace the configured type. + client = _FakeClient({"dssVersion": "14.7.2", "nodeType": "AUTOMATION"}) + monkeypatch.setattr(instances, "get_dss_client", lambda: client) + current = asyncio.run(instances.get_current_instance(FakeContext())) + assert json.loads(current)["instance_type"] == instance_type + assert client.info_calls == 1 + for response in (listing, switching, current): + assert API_KEY not in response + assert "api_key" not in response + + def test_get_current_instance_reports_the_dataiku_version(monkeypatch): client = _FakeClient({"dssVersion": "14.7.2", "nodeType": "DESIGN"}) result = _current(monkeypatch, client) @@ -86,6 +129,7 @@ def missing_client(): assert "dataiku_version" not in result assert result["connection_status"] == "failed" assert result["name"] == "primary" + assert result["instance_type"] == "design" def test_get_current_instance_reports_failed_connection(monkeypatch): diff --git a/tests/test_setup_server.py b/tests/test_setup_server.py index 0721d60f..f0c3f03d 100644 --- a/tests/test_setup_server.py +++ b/tests/test_setup_server.py @@ -72,6 +72,7 @@ def _form(url: str) -> dict[str, list[str]]: "url": [url], "api_key": ["secret"], "description": [""], + "instance_type": ["design"], } @@ -95,6 +96,71 @@ def test_setup_page_explains_instance_url(): assert "Enter the URL of your Dataiku instance." in page +@pytest.mark.parametrize( + "instance_type", ["design", "automation", "deployer", "agent-management"] +) +def test_setup_form_and_page_preserve_instance_type(instance_type): + values = _validate_form( + {**_form("https://example.com"), "instance_type": [instance_type]} + ) + assert values["instance_type"] == instance_type + for page in ( + _page(values=values, connection_validated=True), + _page(values=values, error="Connection failed"), + ): + assert f'' in page + assert '