Skip to content

feat(auth): support Code Studio API ticket authentication - #171

Open
nmartorell wants to merge 10 commits into
mainfrom
feat/code-studios-support
Open

nmartorell wants to merge 10 commits into
mainfrom
feat/code-studios-support

Conversation

@nmartorell

@nmartorell nmartorell commented Sep 30, 2026 •

Copy link
Copy Markdown
Contributor

Purpose

Allow Dataiku Headless to connect automatically to the hosting DSS instance when Headless runs inside a Dataiku Code Studio. The server uses the Code Studio's API ticket and backend URL, so users do not need to configure a personal API key through the setup popup. Explicit environment configuration can still select a different DSS instance.

Changes

  • Load both environment-derived instances: the explicit DKU_DSS_URL / DKU_API_KEY target first, then the hosting Code Studio instance. Select the first at startup while keeping both available through list_instances and switch_instance. The default names are dataiku-from-env and dataiku-from-code-studio; duplicate environment names fail clearly.
  • Integrate Code Studio onboarding and instance-selection precedence into the existing README authentication section, and document the implementation conventions in the coding standards.
  • Support encrypted RPC when DKU_SERVER_CERT is populated. Pass it through the MCP manifest, validate the raw PEM, write it to a private temporary file, and carry its path on the runtime instance. Use that file for HTTPS certificate verification, including hostname checking. Create no certificate files when the variable is absent or empty. Files survive instance switches and configuration resets and are removed on normal process shutdown using atexit, without a global temporary-directory object. Code Studio instances always keep verification enabled; DKU_NO_CHECK_CERTIFICATE applies to the explicit environment instance.
  • Require exactly one non-empty API key or API ticket in stdio configuration and pass tickets to the SDK's internal_ticket parameter. HTTP instances continue using delegated bearer tokens.
  • Exclude API keys, tickets, and certificate paths from current-instance responses. The setup popup remains API-key based.

Validation

  • All 534 automated tests passed, including a local HTTPS server test covering certificate trust and hostname rejection, file permissions and shutdown cleanup, environment precedence and switching, and response redaction.
  • Locked dependency sync, pre-commit checks, and the wheel/source distribution build passed.
  • Live Code Studio testing remains outstanding.

Test in a Code Studio

In a Code Studio template, add an Append to Dockerfile block after the VSCode and Codex blocks, containing:

RUN curl -LsSf https://astral.sh/uv/install.sh | sh

RUN rm -rf /home/dataiku/.codex && codex plugin marketplace add https://github.com/dataiku/dataiku-headless.git --ref feat/code-studios-support
RUN codex plugin add dataiku-headless@dataiku

Build the template and start a Code Studio using it. Open Codex in VSCode and ask it to connect to Dataiku using Dataiku Headless. Without an explicit DKU_DSS_URL, verify that list_instances shows the active instance dataiku-from-code-studio, get_current_instance reports connection_status="connected" with source code-studio-environment, and a read-only call such as list_projects succeeds without entering a personal API key. The current-instance response must contain neither credential nor certificate-path fields.

To test the override, set DKU_DSS_URL and DKU_API_KEY before starting the MCP server. Verify that the explicit instance is selected first and that switch_instance can still select dataiku-from-code-studio. On an encrypted-RPC deployment, repeat the read-only checks with the injected DKU_SERVER_CERT and certificate verification enabled.

@crmapj crmapj left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two inline points on the Code Studio path, plus one docs point:

  • README.md "Auth resolution order" (line 212) still lists only DKU_DSS_URL/DKU_API_KEY and the config file, and line 191 describes env vars as "an explicit override". The Code Studio variables (DKU_IS_CODE_STUDIO, DKU_BACKEND_*, DKU_API_TICKET) and where they sit in that order are currently only documented in CODING_STANDARDS_AND_STRUCTURE.md. AGENTS.md asks for README updates when launch behavior changes.


def _load_instance_from_env_vars() -> DSSInstance | None:
# Code Studio supplies an internal API ticket instead of an API key.
if os.environ.get("DKU_IS_CODE_STUDIO"):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If a user explicitly sets DKU_DSS_URL/DKU_API_KEY inside a Code Studio (e.g. to target a different DSS node), this branch runs first and silently ignores them, so tools act on the hosting instance instead. The README documents these env vars as the explicit override, so letting them win seems safer:

Suggested change
if os.environ.get("DKU_IS_CODE_STUDIO"):
if os.environ.get("DKU_IS_CODE_STUDIO") and not os.environ.get("DKU_DSS_URL"):

f"{os.environ['DKU_BACKEND_PROTOCOL']}://"
f"{os.environ['DKU_BACKEND_HOST']}:{os.environ['DKU_BACKEND_PORT']}"
)
instance = StdioDSSInstanceConfig(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TLS verification when DKU_BACKEND_PROTOCOL=https: Code Studio provides the backend's certificate in DKU_SERVER_CERT, but it isn't read here or forwarded in .mcp.json, and no_check_certificate is always False in this branch. get_dss_client then sets verify=True, so requests only trust the default CA bundle and every call fails with an SSL error unless the backend certificate is publicly trusted. DKU_NO_CHECK_CERTIFICATE isn't honored here either, so there's no workaround. (Not tested against a live HTTPS backend; the https://...:443 test only checks the URL string.)

Rough sketch:

# Code Studio branch: the value may be base64-encoded with a "b64:" prefix
server_cert = os.environ.get("DKU_SERVER_CERT")
if server_cert:
    if server_cert.startswith("b64:"):
        server_cert = base64.b64decode(server_cert[4:]).decode("utf-8")
    # write it to a file and keep the path on the instance, e.g. ca_cert_path

# get_dss_client
client._session.verify = (
    False if current_instance.no_check_certificate
    else current_instance.ca_cert_path or True
)

Plus: add DKU_SERVER_CERT to .mcp.json, and parse DKU_NO_CHECK_CERTIFICATE here the same way as the branch below (a small shared helper would avoid duplicating it).

@nmartorell

nmartorell commented Oct 1, 2026 •

Copy link
Copy Markdown
Contributor Author

@crmapj Thanks for the feedback. The follow-up work addresses the three points:

  1. Explicit override with both instances retained: DKU_DSS_URL / DKU_API_KEY now selects the startup target even inside a Code Studio. The hosting instance is also loaded and remains switchable. Their default names are dataiku-from-env and dataiku-from-code-studio, with an explicit error for duplicate environment names.
  2. README integration: the existing stdio onboarding and authentication section now explains automatic Code Studio connections, explicit overrides, selection order, and switching, alongside API-key setup.
  3. Encrypted RPC: the manifest passes through raw-PEM DKU_SERVER_CERT. When populated, it is validated and written to a private temporary file whose path is used for HTTPS verification. Hostname checking stays enabled. No certificate file is created in the common absent/empty case, and atexit removes files on normal shutdown without a global directory object. Code Studio instances do not consume DKU_NO_CHECK_CERTIFICATE.

All 534 tests pass, including real localhost TLS trust/hostname checks, certificate lifecycle tests, and credential/path redaction. Pre-commit and the package build passed. The PR description and manual Code Studio testing instructions have been updated; live Code Studio testing remains outstanding.

Could you take another look?

Also, I enabled the encrypted RPC functionality of the ai-dev-kit instance, in case you want to test live. Here's project listing from a CS:
image

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants