Add flyte proxy app: authenticated localhost proxy into a Flyte App

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'

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(

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 <free-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 "<unknown>"
+    # 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 "<token>"
+    except Exception:
+        return "<token>"
+
+
+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()

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)