diff --git a/pyproject.toml b/pyproject.toml index 36417d52f..a00c04854 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -261,5 +261,5 @@ ignore-regex = '^\s*"image/\S+": ".*' # ot - variable for output type, nd - NamedDelivery variable, # crasher - legitimate word, tru - intentional truncated JSON in test, # Nome - timezone "America/Nome", lits - short for "literals", -# rouge - NLP scoring metric (not "rogue") -ignore-words-list = 'ot,nd,crasher,tru,nome,lits,rouge' +# rouge - NLP scoring metric (not "rogue"), te - HTTP TE hop-by-hop header +ignore-words-list = 'ot,nd,crasher,tru,nome,lits,rouge,te' diff --git a/src/flyte/app/_runtime/app_serde.py b/src/flyte/app/_runtime/app_serde.py index 24ada3a19..f087f1a38 100644 --- a/src/flyte/app/_runtime/app_serde.py +++ b/src/flyte/app/_runtime/app_serde.py @@ -88,6 +88,7 @@ def _serialized_pod_spec( app_env: AppEnvironment, pod_template: flyte.PodTemplate, serialization_context: SerializationContext, + parameter_overrides: list[Parameter] | None = None, ) -> dict: """ Convert pod spec into a dict for serialization. @@ -144,7 +145,12 @@ def _serialized_pod_spec( if container.name == pod_template.primary_container_name: container.args = app_env.container_args(serialization_context) - container.command = app_env.container_cmd(serialization_context) + # Pass the materialized parameters (delayed values like RunOutput already + # resolved to their File/Dir URI) so the serve container command carries the + # resolved value — matching the container path (get_proto_container). Without + # this the pod path serializes the raw, unresolved parameter and the serve + # entrypoint can't download it. + container.command = app_env.container_cmd(serialization_context, parameter_overrides) limits, requests = {}, {} resources = get_proto_resources(app_env.resources) @@ -184,6 +190,7 @@ def _get_k8s_pod( app_env: AppEnvironment, pod_template: flyte.PodTemplate, serialization_context: SerializationContext, + parameter_overrides: list[Parameter] | None = None, ) -> tasks_pb2.K8sPod: """ Convert pod_template into a K8sPod IDL. @@ -201,7 +208,7 @@ def _get_k8s_pod( from google.protobuf.json_format import Parse from google.protobuf.struct_pb2 import Struct - pod_spec_dict = _serialized_pod_spec(app_env, pod_template, serialization_context) + pod_spec_dict = _serialized_pod_spec(app_env, pod_template, serialization_context, parameter_overrides) pod_spec_idl = Parse(json.dumps(pod_spec_dict), Struct()) metadata = tasks_pb2.K8sObjectMetadata( @@ -396,6 +403,7 @@ async def translate_app_env_to_idl( app_env, app_env.pod_template, serialization_context, + parameter_overrides=parameters, ) elif app_env.image: container = get_proto_container( diff --git a/src/flyte/cli/_proxy.py b/src/flyte/cli/_proxy.py new file mode 100644 index 000000000..c401005a6 --- /dev/null +++ b/src/flyte/cli/_proxy.py @@ -0,0 +1,259 @@ +import asyncio +import errno +import json +import sys + +import rich_click as click + +from . import _common as common + +# Headers that must not be forwarded verbatim across a proxy hop. +_HOP_BY_HOP = { + "connection", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", + "host", + "content-length", +} + + +@click.group(name="proxy") +def proxy(): + """Proxy a local port into a Flyte App through the authenticated app edge.""" + + +@proxy.command(cls=common.CommandBase) +@click.argument("name", type=str, required=False) +@click.option("--url", "url", type=str, default=None, help="Proxy this app URL directly; skip name resolution.") +@click.option("--port", type=int, default=8600, help="Local port to listen on (0 = pick a free port).") +@click.option("--address", type=str, default="127.0.0.1", help="Local bind address; non-loopback triggers a warning.") +@click.option( + "--emit-mcp-config", + is_flag=True, + default=False, + help="Print a generic HTTP-MCP config block for the local endpoint.", +) +@click.option("-v", "--verbose", is_flag=True, default=False, help="Log each proxied request (never the token).") +@click.pass_obj +def app( + cfg: common.CLIConfig, + name: str | None = None, + project: str | None = None, + domain: str | None = None, + url: str | None = None, + port: int = 8600, + address: str = "127.0.0.1", + emit_mcp_config: bool = False, + verbose: bool = False, +): + """ + Open an authenticated localhost proxy into a no-auth Flyte App. + + Reuses the same Union auth (with auto-refresh) the CLI uses, injecting a fresh bearer on every + request, so a local HTTP client — a Grafana/Prometheus MCP, curl, a browser — reaches an + edge-gated app with no token handling. Think: kubectl port-forward for Flyte Apps. + + Foreground; Ctrl-C to stop. + """ + cfg.init(project=project, domain=domain) + + if url: + target = url + elif name: + # Lazy import: keeps `flyte` CLI startup fast — the heavy remote client stack + # only loads when resolving an app by name (the --url path skips it entirely). + import flyte.remote as remote + + target = remote.App.get(name=name).endpoint + else: + raise click.UsageError("Provide an app NAME or --url.") + target = target.rstrip("/") + + label = name or target + try: + asyncio.run(_serve(cfg, target, label, address, port, emit_mcp_config, verbose)) + except KeyboardInterrupt: + pass + + +def _build_authenticator(cfg: common.CLIConfig): + import typing + + from flyte.remote._client.auth._authenticators.factory import get_async_authenticator + from flyte.remote._client.auth._client_config import AuthType, RemoteClientConfigStore + from flyte.remote._client.auth._session import normalize_rpc_endpoint + + plat = cfg.config.platform + if not plat.endpoint: + raise click.UsageError("No endpoint configured; set one via config or FLYTECTL_CONFIG.") + insecure = getattr(plat, "insecure", False) + # Config endpoint is gRPC-style (bare host); the OIDC-metadata client needs an http(s) base. + endpoint = normalize_rpc_endpoint(plat.endpoint, insecure=insecure) + auth_type = typing.cast(AuthType, getattr(plat, "auth_mode", None) or "Pkce") + return get_async_authenticator( + endpoint=endpoint, + cfg_store=RemoteClientConfigStore(endpoint), + auth_type=auth_type, + insecure_skip_verify=getattr(plat, "insecure_skip_verify", False), + ca_cert_file_path=getattr(plat, "ca_cert_file_path", None), + ) + + +def _filter_request_headers(headers) -> dict: + return {k: v for k, v in headers.items() if k.lower() not in _HOP_BY_HOP and k.lower() != "authorization"} + + +def _filter_response_headers(headers) -> dict: + return {k: v for k, v in headers.items() if k.lower() not in _HOP_BY_HOP} + + +async def _serve(cfg, target, label, address, port, emit_mcp_config, verbose): + from aiohttp import ClientSession, ClientTimeout, web + + authenticator = _build_authenticator(cfg) + # Prime credentials (load from keyring / run the auth flow) before we start serving. + if authenticator.get_credentials() is None: + await authenticator.refresh_credentials() + + # auto_decompress=False -> stream upstream bytes verbatim (Content-Encoding preserved). + # total=None -> allow long-lived SSE streams (MCP streamable-HTTP, Grafana Live). + upstream = ClientSession(timeout=ClientTimeout(total=None), auto_decompress=False) + + async def handle(request: "web.Request"): + body = await request.read() + fwd = _filter_request_headers(request.headers) + up_url = target + request.raw_path + + async def send(refresh: bool): + if refresh: + await authenticator.refresh_credentials() + ah = await authenticator.get_auth_headers() + hdrs = dict(fwd) + if ah: + # Normalize the injected bearer onto the standard `authorization` header. + # The SDK auth flow emits its bearer under the CP's configured metadata key + # (typically the gRPC-style `flyte-authorization`), but a Flyte App's HTTP + # auth edge (union-apps) reads only `authorization` and never sees + # `flyte-authorization`, so the request is rejected. Scope this to the + # authenticator's own headers — never rewrite forwarded client headers. + for k, v in ah.headers.items(): + hdrs.pop(k, None) + if isinstance(v, str) and v.startswith("Bearer "): + hdrs["authorization"] = v + else: + hdrs[k] = v + return await upstream.request( + request.method, up_url, headers=hdrs, data=body or None, allow_redirects=False + ) + + resp = await send(refresh=False) + # Stale token: the edge 401s or bounces to /login. Refresh once and retry. + bounced = resp.status in (302, 307) and "/login" in resp.headers.get("Location", "") + if resp.status in (401, 403) or bounced: + resp.release() + resp = await send(refresh=True) + + out = web.StreamResponse(status=resp.status, headers=_filter_response_headers(resp.headers)) + await out.prepare(request) + async for chunk in resp.content.iter_any(): + await out.write(chunk) + await out.write_eof() + resp.release() + if verbose: + click.echo(f"{request.method} {request.path} -> {resp.status}", err=True) + return out + + server = web.Application() + server.router.add_route("*", "/{tail:.*}", handle) + runner = web.AppRunner(server) + await runner.setup() + # Bind loopback on BOTH IPv4 and IPv6 so the proxy owns `localhost` fully. + # macOS resolves `localhost` to ::1 (IPv6) first, and a wildcard listener in + # another process (e.g. OrbStack/Docker on *:PORT) otherwise shadows an + # IPv4-only 127.0.0.1 bind: bind() still succeeds, so the proxy looks healthy + # yet never receives the request (it lands on the other listener, which + # resets it). Binding ::1 too turns that silent half-bind into a loud + # EADDRINUSE. A non-loopback --address is bound as-is (single family). + loopback = address in ("127.0.0.1", "localhost", "::1", "loopback") + hosts = ["127.0.0.1", "::1"] if loopback else [address] + + actual = port + bound: list[str] = [] + for host in hosts: + try: + await web.TCPSite(runner, host, actual).start() + except OSError as e: + if e.errno == errno.EADDRINUSE: + await runner.cleanup() + raise click.ClickException( + f"Port {actual} is already in use (binding {host} failed). Another " + f"process — often a wildcard binder like OrbStack or Docker — holds it " + f"and would silently shadow this proxy. Re-run with --port ." + ) + if host == "::1" and e.errno in (errno.EADDRNOTAVAIL, errno.EAFNOSUPPORT): + click.secho( + f"note: IPv6 loopback unavailable, binding IPv4 only ({e.strerror}).", + fg="yellow", + err=True, + ) + continue + await runner.cleanup() + raise + bound.append(host) + if port == 0 and actual == 0: + # First site picked a free port; pin the other family to the same one. + actual = runner.addresses[0][1] + if not bound: + await runner.cleanup() + raise click.ClickException("Failed to bind any loopback address.") + + display_host = "127.0.0.1" if loopback else address + local = f"http://{display_host}:{actual}" + + if not loopback: + click.secho( + f"WARNING: binding {address} exposes your Union identity to anything that can reach it; prefer 127.0.0.1.", + fg="yellow", + err=True, + ) + identity = _identity(authenticator) + click.echo(f"Proxying {target} -> {local}", err=True) + click.echo(f" authenticating as: {identity}", err=True) + click.echo(" Ctrl-C to stop.", err=True) + if emit_mcp_config: + _emit_mcp_config(label, local) + + try: + await asyncio.Event().wait() # serve until interrupted (Ctrl-C) + finally: + await upstream.close() + await runner.cleanup() + + +def _identity(authenticator) -> str: + creds = authenticator.get_credentials() + if not creds or not creds.access_token: + return "" + # Best-effort: decode the JWT payload's sub/email without verifying (display only). + try: + import base64 + + payload = creds.access_token.split(".")[1] + payload += "=" * (-len(payload) % 4) + claims = json.loads(base64.urlsafe_b64decode(payload)) + return claims.get("email") or claims.get("sub") or "" + except Exception: + return "" + + +def _emit_mcp_config(name: str, local: str): + block = {"mcpServers": {name: {"type": "http", "url": local}}} + click.echo("", err=True) + click.echo("# MCP client config (generic HTTP transport) — point your client at the local proxy:", err=True) + sys.stdout.write(json.dumps(block, indent=2) + "\n") + sys.stdout.flush() diff --git a/src/flyte/cli/main.py b/src/flyte/cli/main.py index 5b1d0cf5e..9f6d82abf 100644 --- a/src/flyte/cli/main.py +++ b/src/flyte/cli/main.py @@ -16,6 +16,7 @@ from ._get import get from ._plugins import discover_and_register_plugins from ._prefetch import prefetch +from ._proxy import proxy from ._rerun import rerun from ._run import run from ._serve import serve @@ -303,6 +304,7 @@ def main( main.add_command(stop) # type: ignore main.add_command(prefetch) # type: ignore main.add_command(edit) # type: ignore +main.add_command(proxy) # type: ignore # Discover and register CLI plugins from installed packages discover_and_register_plugins(main) diff --git a/tests/cli/test_proxy.py b/tests/cli/test_proxy.py new file mode 100644 index 000000000..107975cca --- /dev/null +++ b/tests/cli/test_proxy.py @@ -0,0 +1,75 @@ +"""Unit tests for flyte.cli._proxy (the `flyte proxy app` command helpers).""" + +from __future__ import annotations + +import base64 +import json +from unittest.mock import MagicMock + +from flyte.cli._proxy import ( + _emit_mcp_config, + _filter_request_headers, + _filter_response_headers, + _identity, +) + + +def _jwt(claims: dict) -> str: + def b64(obj: dict) -> str: + return base64.urlsafe_b64encode(json.dumps(obj).encode()).rstrip(b"=").decode() + + return f"{b64({'alg': 'none'})}.{b64(claims)}.signature" + + +class TestHeaderFiltering: + def test_request_headers_drop_authorization_and_hop_by_hop(self): + out = _filter_request_headers( + { + "Authorization": "Bearer inbound", + "Host": "localhost:8600", + "Connection": "keep-alive", + "Content-Length": "10", + "Accept": "*/*", + "X-Custom": "keep", + } + ) + lowered = {k.lower() for k in out} + assert "authorization" not in lowered # inbound auth must never be forwarded + assert "host" not in lowered + assert "connection" not in lowered + assert "content-length" not in lowered + assert out["Accept"] == "*/*" + assert out["X-Custom"] == "keep" + + def test_response_headers_drop_hop_by_hop_keep_others(self): + out = _filter_response_headers( + {"Transfer-Encoding": "chunked", "Content-Type": "application/json", "Content-Encoding": "gzip"} + ) + lowered = {k.lower() for k in out} + assert "transfer-encoding" not in lowered + assert out["Content-Type"] == "application/json" + assert out["Content-Encoding"] == "gzip" # preserved for byte-verbatim streaming + + +class TestIdentity: + def test_prefers_email_claim(self): + auth = MagicMock() + auth.get_credentials.return_value = MagicMock(access_token=_jwt({"email": "me@union.ai", "sub": "abc"})) + assert _identity(auth) == "me@union.ai" + + def test_falls_back_to_sub(self): + auth = MagicMock() + auth.get_credentials.return_value = MagicMock(access_token=_jwt({"sub": "subject-123"})) + assert _identity(auth) == "subject-123" + + def test_unknown_when_no_credentials(self): + auth = MagicMock() + auth.get_credentials.return_value = None + assert _identity(auth) == "" + + +class TestEmitMcpConfig: + def test_emits_generic_http_block(self, capsys): + _emit_mcp_config("grafana", "http://127.0.0.1:8600") + block = json.loads(capsys.readouterr().out) + assert block == {"mcpServers": {"grafana": {"type": "http", "url": "http://127.0.0.1:8600"}}}