feat(auth): support Code Studio API ticket authentication - #171
nmartorell wants to merge 10 commits into
Conversation
crmapj
left a comment
There was a problem hiding this comment.
Two inline points on the Code Studio path, plus one docs point:
README.md"Auth resolution order" (line 212) still lists onlyDKU_DSS_URL/DKU_API_KEYand 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 inCODING_STANDARDS_AND_STRUCTURE.md.AGENTS.mdasks 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"): |
There was a problem hiding this comment.
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:
| 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( |
There was a problem hiding this comment.
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).
|
@crmapj Thanks for the feedback. The follow-up work addresses the three points:
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: |

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
DKU_DSS_URL/DKU_API_KEYtarget first, then the hosting Code Studio instance. Select the first at startup while keeping both available throughlist_instancesandswitch_instance. The default names aredataiku-from-envanddataiku-from-code-studio; duplicate environment names fail clearly.DKU_SERVER_CERTis 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 usingatexit, without a global temporary-directory object. Code Studio instances always keep verification enabled;DKU_NO_CHECK_CERTIFICATEapplies to the explicit environment instance.internal_ticketparameter. HTTP instances continue using delegated bearer tokens.Validation
Test in a Code Studio
In a Code Studio template, add an Append to Dockerfile block after the VSCode and Codex blocks, containing:
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 thatlist_instancesshows the active instancedataiku-from-code-studio,get_current_instancereportsconnection_status="connected"with sourcecode-studio-environment, and a read-only call such aslist_projectssucceeds 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_URLandDKU_API_KEYbefore starting the MCP server. Verify that the explicit instance is selected first and thatswitch_instancecan still selectdataiku-from-code-studio. On an encrypted-RPC deployment, repeat the read-only checks with the injectedDKU_SERVER_CERTand certificate verification enabled.