mlspace-plugin 0.2.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (53) hide show
  1. mlspace_mcp/__init__.py +3 -0
  2. mlspace_mcp/__main__.py +125 -0
  3. mlspace_mcp/_status.py +18 -0
  4. mlspace_mcp/auth.py +109 -0
  5. mlspace_mcp/bodyspec.py +184 -0
  6. mlspace_mcp/client.py +170 -0
  7. mlspace_mcp/client_setup.py +235 -0
  8. mlspace_mcp/config.py +283 -0
  9. mlspace_mcp/errors.py +241 -0
  10. mlspace_mcp/formatting.py +502 -0
  11. mlspace_mcp/init_cli.py +356 -0
  12. mlspace_mcp/job_hooks.py +242 -0
  13. mlspace_mcp/jsonc.py +90 -0
  14. mlspace_mcp/overview.py +421 -0
  15. mlspace_mcp/registry.py +632 -0
  16. mlspace_mcp/resource_refs.py +188 -0
  17. mlspace_mcp/server.py +190 -0
  18. mlspace_mcp/setup_ui.py +103 -0
  19. mlspace_mcp/skills/mlspace-call-deployed-model-predict/SKILL.md +36 -0
  20. mlspace_mcp/skills/mlspace-deploy-inference-service/SKILL.md +42 -0
  21. mlspace_mcp/skills/mlspace-diagnose-data-transfer/SKILL.md +41 -0
  22. mlspace_mcp/skills/mlspace-diagnose-failed-training-job/SKILL.md +64 -0
  23. mlspace_mcp/skills/mlspace-diagnose-job-launch-rejection/SKILL.md +44 -0
  24. mlspace_mcp/skills/mlspace-diagnose-unhealthy-inference/SKILL.md +40 -0
  25. mlspace_mcp/skills/mlspace-discover-launchable-resources/SKILL.md +45 -0
  26. mlspace_mcp/skills/mlspace-launch-training-job/SKILL.md +66 -0
  27. mlspace_mcp/skills/mlspace-list-running-jobs/SKILL.md +43 -0
  28. mlspace_mcp/skills/mlspace-resume-notebook-and-set-autoshutdown/SKILL.md +44 -0
  29. mlspace_mcp/skills/mlspace-stage-s3-to-nfs/SKILL.md +45 -0
  30. mlspace_mcp/skills/mlspace-transfer-compute-node/SKILL.md +63 -0
  31. mlspace_mcp/spec/body_schemas.json +7506 -0
  32. mlspace_mcp/spec/openapi.json +28347 -0
  33. mlspace_mcp/tls.py +46 -0
  34. mlspace_mcp/tools/__init__.py +1 -0
  35. mlspace_mcp/tools/allocations.py +92 -0
  36. mlspace_mcp/tools/async_inference.py +77 -0
  37. mlspace_mcp/tools/build_image.py +76 -0
  38. mlspace_mcp/tools/dalle.py +52 -0
  39. mlspace_mcp/tools/data_transfer.py +260 -0
  40. mlspace_mcp/tools/docker_registry.py +156 -0
  41. mlspace_mcp/tools/inference.py +88 -0
  42. mlspace_mcp/tools/jobs.py +155 -0
  43. mlspace_mcp/tools/notebooks.py +174 -0
  44. mlspace_mcp/tools/queues.py +202 -0
  45. mlspace_mcp/tools/resources.py +86 -0
  46. mlspace_mcp/tools/tensorboards.py +112 -0
  47. mlspace_mcp/tools/workspaces.py +92 -0
  48. mlspace_mcp/workspace_context.py +107 -0
  49. mlspace_plugin-0.2.1.dist-info/METADATA +427 -0
  50. mlspace_plugin-0.2.1.dist-info/RECORD +53 -0
  51. mlspace_plugin-0.2.1.dist-info/WHEEL +4 -0
  52. mlspace_plugin-0.2.1.dist-info/entry_points.txt +2 -0
  53. mlspace_plugin-0.2.1.dist-info/licenses/LICENSE +21 -0
@@ -0,0 +1,3 @@
1
+ """mlspace_mcp — MCP server for the Cloud.ru MLSpace public API v2."""
2
+
3
+ __version__ = "0.2.1"
@@ -0,0 +1,125 @@
1
+ """CLI entry point for mlspace-plugin."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import argparse
6
+ import logging
7
+ import os
8
+ import sys
9
+ from pathlib import Path
10
+
11
+ from . import __version__
12
+ from .config import load_settings, user_config_file
13
+ from .server import build_server
14
+
15
+
16
+ def _configure_logging() -> None:
17
+ """Basic stderr logging so the server's logger.warning/exception emit in prod.
18
+ Level via MLSPACE_LOG_LEVEL (default INFO). stderr keeps stdout clean for stdio MCP.
19
+ """
20
+ level = os.environ.get("MLSPACE_LOG_LEVEL", "INFO").upper()
21
+ logging.basicConfig(
22
+ level=getattr(logging, level, logging.INFO),
23
+ format="%(asctime)s %(levelname)s %(name)s: %(message)s",
24
+ stream=sys.stderr,
25
+ )
26
+
27
+
28
+ def main(argv: list[str] | None = None) -> int:
29
+ parser = argparse.ArgumentParser(
30
+ prog="mlspace-plugin",
31
+ description="MCP server for the Cloud.ru MLSpace public API v2.",
32
+ )
33
+ parser.add_argument("--version", action="version", version=f"mlspace-plugin {__version__}")
34
+ parser.add_argument(
35
+ "--transport",
36
+ choices=["streamable-http", "stdio"],
37
+ help="Override MLSPACE_TRANSPORT.",
38
+ )
39
+ parser.add_argument(
40
+ "--list-tools",
41
+ action="store_true",
42
+ help="Print the tools that would be registered, then exit (no credentials needed).",
43
+ )
44
+ # Optional subcommand: with none given the CLI serves, as it always has.
45
+ sub = parser.add_subparsers(dest="command")
46
+ setup = sub.add_parser(
47
+ "setup",
48
+ help="Set up access keys, workspaces, MCP clients and skills.",
49
+ description=(
50
+ "Ask for Cloud.ru Key ID and Key Secret, discover MLSpace workspaces, and write them to "
51
+ f"{user_config_file()} with mode 600. Runs in your terminal: the values "
52
+ "are never sent to a model."
53
+ ),
54
+ )
55
+ setup.add_argument("--path", type=Path, help="Write somewhere else than the user config.")
56
+ setup.add_argument("--force", action="store_true", help="Overwrite an existing file.")
57
+ setup.add_argument(
58
+ "--base-url",
59
+ help="MLSpace API endpoint to save (defaults to production or the existing target).",
60
+ )
61
+ setup.add_argument(
62
+ "--no-verify", action="store_true", help="Save unverified config offline; requires --config-only and --workspace."
63
+ )
64
+ setup.add_argument("--client", dest="clients", action="append", choices=["claude-code", "codex", "opencode"], help="Connect this client (repeat for several); default: auto-detect.")
65
+ setup.add_argument("--client-path", dest="client_paths", action="append", metavar="CLIENT=PATH", help="Explicit client executable when not in PATH.")
66
+ setup.add_argument("--config-only", action="store_true", help="Only configure credentials/workspaces; do not register clients or skills.")
67
+ setup.add_argument("--non-interactive", action="store_true", help="Never prompt; require explicit workspace IDs and credential source or saved keys.")
68
+ setup.add_argument("--from-env", action="store_true", help="Use MLSPACE_CLIENT_ID / MLSPACE_CLIENT_SECRET from the environment.")
69
+ setup.add_argument("--credentials-file", type=Path, help="Import keys from an existing dotenv file without printing them.")
70
+ setup.add_argument("--replace-credentials", action="store_true", help="Ask for new keys instead of reusing saved keys.")
71
+ setup.add_argument("--workspace", dest="workspace_ids", action="append", help="Exact workspace ID to select (repeat for several).")
72
+ setup.add_argument("--ca-file", type=Path, help="Trusted corporate CA PEM; a copy is saved for subsequent MCP launches.")
73
+ setup.add_argument("--list-workspaces", action="store_true", help="Return a bounded JSON list without saving or installing anything.")
74
+ setup.add_argument("--search", default="", help="Filter workspace JSON by name, project or ID.")
75
+ setup.add_argument("--offset", type=int, default=0)
76
+ setup.add_argument("--limit", type=int, default=20, help="Workspace JSON page size (1–100).")
77
+ args = parser.parse_args(argv)
78
+
79
+ if args.command == "setup":
80
+ from .init_cli import run_init
81
+
82
+ return run_init(
83
+ args.path,
84
+ force=args.force,
85
+ verify=not args.no_verify,
86
+ base_url=args.base_url,
87
+ config_only=args.config_only, non_interactive=args.non_interactive,
88
+ from_env=args.from_env, credentials_file=args.credentials_file,
89
+ workspace_ids=args.workspace_ids, clients=args.clients, client_paths=args.client_paths,
90
+ ca_file=args.ca_file, replace_credentials=args.replace_credentials,
91
+ list_workspaces=args.list_workspaces, search=args.search, offset=args.offset, limit=args.limit,
92
+ )
93
+
94
+ _configure_logging()
95
+ settings = load_settings()
96
+ if args.transport:
97
+ settings.transport = args.transport
98
+
99
+ if args.list_tools:
100
+ mcp = build_server(settings)
101
+ tools = mcp._tool_manager.list_tools()
102
+ mode = "read-only" if settings.readonly else "read-write"
103
+ print(f"mlspace-plugin {__version__} — {len(tools)} tools ({mode}):")
104
+ for tool in sorted(tools, key=lambda t: t.name):
105
+ print(f" {tool.name}")
106
+ return 0
107
+
108
+ settings.validate_bind()
109
+ settings.require_credentials()
110
+ mcp = build_server(settings)
111
+
112
+ if settings.transport == "stdio":
113
+ mcp.run(transport="stdio")
114
+ else:
115
+ print(
116
+ f"mlspace-plugin serving streamable-HTTP on "
117
+ f"http://{settings.host}:{settings.port}/mcp",
118
+ file=sys.stderr,
119
+ )
120
+ mcp.run(transport="streamable-http")
121
+ return 0
122
+
123
+
124
+ if __name__ == "__main__":
125
+ raise SystemExit(main())
mlspace_mcp/_status.py ADDED
@@ -0,0 +1,18 @@
1
+ """Single source of truth for the JobStatus canonical spelling map.
2
+
3
+ Kept in its own tiny leaf module (it imports NOTHING from the package) so both
4
+ ``registry.py`` and ``formatting.py`` can share ONE copy without reintroducing a
5
+ tools -> registry -> formatting import cycle. Previously each held an intentional
6
+ twin copy guarded by a drift test; now there is exactly one.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ # Mirrors the spec ``JobStatus-Input`` enum: casefold -> canonical spelling.
12
+ JOB_STATUS_CANON: dict[str, str] = {
13
+ s.casefold(): s
14
+ for s in (
15
+ "Completed", "Completing", "Deleted", "Failed", "Pending", "Running",
16
+ "Stopped", "Succeeded", "Terminated", "Terminating", "Restarting",
17
+ )
18
+ }
mlspace_mcp/auth.py ADDED
@@ -0,0 +1,109 @@
1
+ """Service-account token management for the MLSpace API.
2
+
3
+ ``POST /public/v2/service_auth`` exchanges ``{client_id, client_secret}`` for a
4
+ short-lived (≈1h) bearer token. The response nests the token:
5
+ ``{"token": {"access_token": "...", "expires_in": 3600}}`` — the sibling
6
+ status/error_* fields are deprecated and must NOT be read.
7
+
8
+ Refresh is single-flight: concurrent callers (and concurrent 401 retries against
9
+ a revoked token) collapse into exactly one ``/service_auth`` call.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ import asyncio
15
+ import time
16
+
17
+ import httpx
18
+
19
+ from .errors import MLSpaceError, error_from_exception, error_from_response
20
+
21
+ _EXPIRY_SKEW = 60.0 # refresh this many seconds before the token actually expires
22
+
23
+
24
+ class TokenManager:
25
+ def __init__(
26
+ self,
27
+ http: httpx.AsyncClient,
28
+ base_url: str,
29
+ client_id: str,
30
+ client_secret: str,
31
+ ):
32
+ self._http = http
33
+ self._url = base_url.rstrip("/") + "/public/v2/service_auth"
34
+ self._client_id = client_id
35
+ self._client_secret = client_secret
36
+ self._token: str | None = None
37
+ self._expiry: float = 0.0
38
+ self._version: int = 0
39
+ self._lock = asyncio.Lock()
40
+
41
+ @staticmethod
42
+ def _now() -> float:
43
+ return time.monotonic()
44
+
45
+ def _valid(self) -> bool:
46
+ return self._token is not None and self._now() < self._expiry
47
+
48
+ async def get_token(self) -> tuple[str, int]:
49
+ """Return ``(token, version)``, refreshing if missing/expired."""
50
+ if self._valid():
51
+ return self._token, self._version # type: ignore[return-value]
52
+ async with self._lock:
53
+ if self._valid(): # someone refreshed while we waited
54
+ return self._token, self._version # type: ignore[return-value]
55
+ await self._refresh()
56
+ return self._token, self._version # type: ignore[return-value]
57
+
58
+ async def refresh_after_401(self, used_version: int) -> tuple[str, int]:
59
+ """Force a refresh after a 401, but only if no one else already has.
60
+
61
+ All concurrent 401s share the same ``used_version`` so they collapse into
62
+ a single ``/service_auth`` call.
63
+ """
64
+ async with self._lock:
65
+ if self._version != used_version and self._valid():
66
+ return self._token, self._version # type: ignore[return-value]
67
+ await self._refresh()
68
+ return self._token, self._version # type: ignore[return-value]
69
+
70
+ async def _refresh(self) -> None:
71
+ # IMPORTANT: never send Authorization/x-api-key/x-workspace-id here.
72
+ try:
73
+ resp = await self._http.post(
74
+ self._url,
75
+ json={"client_id": self._client_id, "client_secret": self._client_secret},
76
+ )
77
+ except Exception as exc: # transport-level
78
+ raise error_from_exception(exc) from exc
79
+
80
+ if resp.status_code == 401:
81
+ raise MLSpaceError(
82
+ "Invalid service-account credentials (401 from /service_auth).",
83
+ status=401,
84
+ hint="Check MLSPACE_CLIENT_ID and MLSPACE_CLIENT_SECRET.",
85
+ )
86
+ if resp.status_code >= 400:
87
+ raise error_from_response(resp)
88
+
89
+ try:
90
+ payload = resp.json()
91
+ except Exception as exc:
92
+ raise MLSpaceError("Auth response was not valid JSON.") from exc
93
+
94
+ token_obj = (payload or {}).get("token") or {}
95
+ access_token = token_obj.get("access_token")
96
+ expires_in = token_obj.get("expires_in")
97
+ if not access_token:
98
+ raise MLSpaceError(
99
+ "Auth response missing token.access_token.",
100
+ hint="Unexpected /service_auth response shape.",
101
+ )
102
+ try:
103
+ ttl = float(expires_in) if expires_in is not None else 3600.0
104
+ except (TypeError, ValueError):
105
+ ttl = 3600.0
106
+ self._token = access_token
107
+ # non-positive TTL (0 / negative) -> treat as already expired, refresh next call
108
+ self._expiry = self._now() + (max(ttl - _EXPIRY_SKEW, 1.0) if ttl > 0 else 0.0)
109
+ self._version += 1
@@ -0,0 +1,184 @@
1
+ """Spec-driven request-body help + LENIENT validation.
2
+
3
+ Schemas come from the bundled build artifact ``spec/body_schemas.json`` (one
4
+ self-contained JSON Schema per "METHOD /path", refs under ``$defs``). Loaded once.
5
+
6
+ Design stance (deliberate — see DESIGN.md): the live API is the real validator.
7
+ ``lint`` is **fail-open** — it reports only HIGH-CONFIDENCE structural problems
8
+ (missing required field incl. nested, clear scalar-type mismatch, enum
9
+ violation) and stays silent on anything it can't be sure about (anyOf/oneOf
10
+ variants, unknown constructs, extra properties). It must never reject a body the
11
+ API would accept; the goal is to catch the obvious "forgot a required field"
12
+ before a round-trip, not to re-implement OpenAPI validation.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import json
18
+ from functools import lru_cache
19
+ from pathlib import Path
20
+ from typing import Any
21
+
22
+ _ARTIFACT = Path(__file__).parent / "spec" / "body_schemas.json"
23
+ _MAX_DEPTH = 8
24
+
25
+
26
+ @lru_cache(maxsize=1)
27
+ def _all() -> dict[str, dict]:
28
+ if not _ARTIFACT.exists():
29
+ return {}
30
+ return json.loads(_ARTIFACT.read_text())
31
+
32
+
33
+ def schema_for(method: str, path: str) -> dict | None:
34
+ return _all().get(f"{method.upper()} {path}")
35
+
36
+
37
+ def _resolve(node: Any, defs: dict) -> Any:
38
+ """Resolve a $ref / single-ref allOf to its target schema (one hop)."""
39
+ seen = 0
40
+ while isinstance(node, dict) and seen < 16:
41
+ if "$ref" in node:
42
+ name = node["$ref"].split("/")[-1]
43
+ node = defs.get(name, {})
44
+ elif "allOf" in node and len(node["allOf"]) == 1 and "$ref" in node["allOf"][0]:
45
+ name = node["allOf"][0]["$ref"].split("/")[-1]
46
+ target = defs.get(name, {})
47
+ # merge sibling keys (e.g. description) over the target
48
+ node = {**target, **{k: v for k, v in node.items() if k != "allOf"}}
49
+ else:
50
+ break
51
+ seen += 1
52
+ return node
53
+
54
+
55
+ def _types(schema: dict) -> set[str]:
56
+ t = schema.get("type")
57
+ if isinstance(t, str):
58
+ return {t}
59
+ if isinstance(t, list):
60
+ return set(t)
61
+ return set()
62
+
63
+
64
+ def _enum_str(schema: dict) -> str:
65
+ e = schema.get("enum")
66
+ if not isinstance(e, list):
67
+ return ""
68
+ if len(e) <= 6:
69
+ return ":" + "|".join(str(x) for x in e)
70
+ return f":<{len(e)} options>" # mark long enums instead of dropping silently
71
+
72
+
73
+ # ---- description hint -------------------------------------------------------
74
+
75
+
76
+ def summarize(schema: dict) -> str:
77
+ """One-line shape hint for the tool description (required fields + nested)."""
78
+ defs = schema.get("$defs", {})
79
+ root = _resolve(schema, defs)
80
+ if not isinstance(root, dict) or "properties" not in root:
81
+ return ""
82
+ required = root.get("required", []) or []
83
+ props = root.get("properties", {})
84
+
85
+ def field(name: str) -> str:
86
+ sub = _resolve(props.get(name, {}), defs)
87
+ if "object" in _types(sub) or "properties" in sub:
88
+ sub_req = sub.get("required", []) or []
89
+ inner = ", ".join(
90
+ n + _enum_str(_resolve(sub.get("properties", {}).get(n, {}), defs))
91
+ for n in sub_req
92
+ )
93
+ return f"{name}{{{inner}}}" if inner else f"{name}{{...}}"
94
+ return name + _enum_str(sub)
95
+
96
+ req = [field(n) for n in required]
97
+ optional_all = [n for n in props if n not in required]
98
+ optional = optional_all[:6]
99
+ parts = []
100
+ if req:
101
+ parts.append("required: " + ", ".join(req))
102
+ if optional:
103
+ more = f" (+{len(optional_all) - len(optional)} more)" if len(optional_all) > len(optional) else ""
104
+ parts.append("optional: " + ", ".join(optional) + more)
105
+ return "; ".join(parts)
106
+
107
+
108
+ # ---- lenient validation -----------------------------------------------------
109
+
110
+ _JSON_TYPE: dict[str, Any] = {
111
+ "string": str,
112
+ "integer": int,
113
+ "number": (int, float),
114
+ "boolean": bool,
115
+ "array": list,
116
+ "object": dict,
117
+ }
118
+
119
+
120
+ def _type_ok(value: Any, types: set[str]) -> bool:
121
+ if not types:
122
+ return True
123
+ if value is None:
124
+ return "null" in types or True # nullability is often implicit; stay lenient
125
+ for t in types:
126
+ py = _JSON_TYPE.get(t)
127
+ if py is None:
128
+ return True # unknown type token -> don't judge
129
+ if t in ("integer", "number") and isinstance(value, bool):
130
+ continue # bool is not a number here
131
+ if t == "integer" and isinstance(value, float) and value.is_integer():
132
+ return True # 5.0 for an integer field: backends coerce — don't reject
133
+ if isinstance(value, py):
134
+ return True
135
+ return False
136
+
137
+
138
+ def lint(schema: dict, value: Any) -> list[str]:
139
+ """Return high-confidence problems with ``value`` against ``schema`` (or [])."""
140
+ defs = schema.get("$defs", {})
141
+ errors: list[str] = []
142
+
143
+ def check(node: Any, val: Any, path: str, depth: int) -> None:
144
+ if depth > _MAX_DEPTH:
145
+ return
146
+ node = _resolve(node, defs)
147
+ if not isinstance(node, dict):
148
+ return
149
+ # unions / combinators: too ambiguous to enforce safely -> skip
150
+ if any(k in node for k in ("anyOf", "oneOf")) or (
151
+ "allOf" in node and len(node.get("allOf", [])) > 1
152
+ ):
153
+ return
154
+
155
+ types = _types(node)
156
+ if types and not _type_ok(val, types):
157
+ errors.append(f"{path or 'body'}: expected {'/'.join(sorted(types))}, got {type(val).__name__}")
158
+ return
159
+
160
+ # NOTE: enum membership is deliberately NOT enforced here. Enums are exactly
161
+ # what a cloud platform expands over time (new regions/connectors/...); a
162
+ # hard local check would fail-closed on a value the live API now accepts,
163
+ # violating this module's fail-open contract. The API (422) validates enums;
164
+ # summarize() still surfaces the known values as a hint to the model.
165
+
166
+ if ("object" in types or "properties" in node) and isinstance(val, dict):
167
+ for req in node.get("required", []) or []:
168
+ # flag only when the KEY is absent; explicit null is left to the API
169
+ # (a required+nullable field sent as null is valid — never reject it)
170
+ if req not in val:
171
+ errors.append(f"{path + '.' if path else ''}{req}: required field is missing")
172
+ for key, sub in (node.get("properties") or {}).items():
173
+ if key in val and val[key] is not None:
174
+ check(sub, val[key], f"{path + '.' if path else ''}{key}", depth + 1)
175
+ elif "array" in types and isinstance(val, list):
176
+ items = node.get("items")
177
+ if isinstance(items, dict) and val:
178
+ check(items, val[0], f"{path}[]", depth + 1)
179
+
180
+ try:
181
+ check(schema, value, "", 0)
182
+ except Exception:
183
+ return [] # fail-open: never block on an internal hiccup
184
+ return errors
mlspace_mcp/client.py ADDED
@@ -0,0 +1,170 @@
1
+ """The single HTTP entry point to the MLSpace API.
2
+
3
+ The client owns ALL header injection (Authorization + x-api-key + x-workspace-id),
4
+ path templating, query encoding (repeated keys for list values), and JSON bodies
5
+ (attached for any method, including DELETE). Tools never touch headers.
6
+
7
+ Frozen contract (do not change without updating every tools/*.py):
8
+
9
+ async def request(
10
+ self,
11
+ method: str,
12
+ path_template: str,
13
+ *,
14
+ path_params: dict[str, str] | None = None,
15
+ params: dict | list[tuple] | None = None,
16
+ json_body: Any | None = None,
17
+ ) -> Any
18
+ """
19
+
20
+ from __future__ import annotations
21
+
22
+ import asyncio
23
+ import re
24
+ from typing import Any
25
+
26
+ import httpx
27
+
28
+ from .auth import TokenManager
29
+ from .errors import MLSpaceError, error_from_exception, error_from_response
30
+
31
+ # Defence-in-depth: even values that reach the client WITHOUT going through
32
+ # resolve_request's per-segment encoding (the guard/namespace prechecks call
33
+ # ``request`` directly with a raw id) must never carry a traversal or control char
34
+ # into the final path. resolve_request is the primary route-confusion guard; this
35
+ # is the last line before the wire.
36
+ _CTRL_RE = re.compile(r"[\x00-\x1f\x7f]")
37
+
38
+
39
+ def _assert_safe_path(path: str) -> None:
40
+ if _CTRL_RE.search(path):
41
+ raise MLSpaceError("Refusing to send a request whose path contains control chars.")
42
+ for seg in path.split("/"):
43
+ if seg in ("..", "."):
44
+ raise MLSpaceError(
45
+ f"Refusing a request whose path contains a '{seg}' traversal segment "
46
+ f"({path!r}) — a path parameter must be a single opaque id/name."
47
+ )
48
+
49
+
50
+ class MLSpaceClient:
51
+ def __init__(
52
+ self,
53
+ http: httpx.AsyncClient,
54
+ tokens: TokenManager,
55
+ *,
56
+ base_url: str,
57
+ api_key: str,
58
+ workspace_id: str,
59
+ namespace: str = "",
60
+ ):
61
+ self._http = http
62
+ self._tokens = tokens
63
+ self._base = base_url.rstrip("/")
64
+ self._api_key = api_key
65
+ self._workspace_id = workspace_id
66
+ # k8s namespace of the configured workspace; resolved lazily from
67
+ # workspace_id when not pinned via config.
68
+ self._namespace: str | None = namespace or None
69
+ self._ns_lock = asyncio.Lock()
70
+
71
+ @property
72
+ def workspace_id(self) -> str:
73
+ return self._workspace_id
74
+
75
+ async def workspace_namespace(self) -> str | None:
76
+ """The configured workspace's k8s namespace, resolved once and cached.
77
+
78
+ Returns None if it can't be determined (caller then falls back to asking
79
+ the model for it). Single-flight: concurrent first-callers collapse into
80
+ one ``workspaces/v3/{id}`` lookup.
81
+ """
82
+ if self._namespace is not None:
83
+ return self._namespace
84
+ async with self._ns_lock:
85
+ if self._namespace is not None:
86
+ return self._namespace
87
+ try:
88
+ data = await self.request(
89
+ "GET",
90
+ "/public/v2/workspaces/v3/{workspace_id}",
91
+ path_params={"workspace_id": self._workspace_id},
92
+ )
93
+ except Exception:
94
+ return None # leave uncached: retry on a later call
95
+ ns = data.get("namespace") if isinstance(data, dict) else None
96
+ if ns:
97
+ self._namespace = ns
98
+ return ns
99
+
100
+ def _headers(self, token: str) -> dict[str, str]:
101
+ return {
102
+ "Authorization": f"Bearer {token}",
103
+ "x-api-key": self._api_key,
104
+ "x-workspace-id": self._workspace_id,
105
+ }
106
+
107
+ async def request(
108
+ self,
109
+ method: str,
110
+ path_template: str,
111
+ *,
112
+ path_params: dict[str, str] | None = None,
113
+ params: Any = None,
114
+ json_body: Any = None,
115
+ ) -> Any:
116
+ if path_params:
117
+ try:
118
+ path = path_template.format(**path_params)
119
+ except KeyError as exc:
120
+ raise MLSpaceError(
121
+ f"Missing path parameter {exc} for {path_template}."
122
+ ) from exc
123
+ else:
124
+ path = path_template
125
+ _assert_safe_path(path)
126
+ url = self._base + path
127
+
128
+ token, version = await self._tokens.get_token()
129
+ resp = await self._send(method, url, token, params, json_body)
130
+
131
+ if resp.status_code == 401:
132
+ # token may be revoked/expired-early — single-flight force refresh, retry once
133
+ token, _ = await self._tokens.refresh_after_401(version)
134
+ resp = await self._send(method, url, token, params, json_body)
135
+
136
+ if resp.status_code >= 400:
137
+ raise error_from_response(resp)
138
+ return self._parse(resp)
139
+
140
+ async def _send(
141
+ self, method: str, url: str, token: str, params: Any, json_body: Any
142
+ ) -> httpx.Response:
143
+ kwargs: dict[str, Any] = {"headers": self._headers(token)}
144
+ if params is not None:
145
+ kwargs["params"] = params
146
+ if json_body is not None:
147
+ kwargs["json"] = json_body # attached for any method, incl. DELETE
148
+ try:
149
+ return await self._http.request(method, url, **kwargs)
150
+ except MLSpaceError:
151
+ raise
152
+ except Exception as exc:
153
+ raise error_from_exception(exc) from exc
154
+
155
+ @staticmethod
156
+ def _parse(resp: httpx.Response) -> Any:
157
+ if resp.status_code == 204 or not resp.content:
158
+ return None
159
+ ctype = resp.headers.get("content-type", "").lower()
160
+ # explicit binary (e.g. dalle image) — don't dump raw bytes into context
161
+ if ctype.startswith(("image/", "audio/", "video/", "application/octet-stream")):
162
+ return f"<binary response, {len(resp.content)} bytes, content-type={ctype}>"
163
+ # json / text / unknown-or-missing content-type — try json, then text
164
+ try:
165
+ return resp.json()
166
+ except Exception:
167
+ try:
168
+ return resp.text
169
+ except Exception:
170
+ return f"<{len(resp.content)} bytes, content-type={ctype or 'unknown'}>"