Union MCP

Instructions

- Based on the prompt and inputs dictionary, determine the task to run
- Format the inputs dictionary so that it matches the task function signature
- Invoke the task

Args:
    project: Project to run the task in.
    domain: Domain to run the task in.
    name: Name of the task to run.
    inputs: A dictionary of inputs to the task.

Returns:
    A dictionary of outputs from the task.

Implementation Reference

• union_mcp/v2/server.py:38-66 (handler)

  Primary MCP tool handler and registration for 'run_task' in v2 server. Decorated with @mcp.tool() and @require_auth. Initializes Flyte and calls the resources helper.

  @mcp.tool()
  @require_auth
  async def run_task(
      name: str,
      inputs: dict,
      project: str,
      domain: str,
      ctx: Context,
  ) -> dict:
      ctx.info(f"Running task {name} in project {project} and domain {domain}")
      """Run a task with natural language.
  
      - Based on the prompt and inputs dictionary, determine the task to run
      - Format the inputs dictionary so that it matches the task function signature
      - Invoke the task
      
      Args:
          project: Project to run the task in.
          domain: Domain to run the task in.
          name: Name of the task to run.
          inputs: A dictionary of inputs to the task.
  
      Returns:
          A dictionary of outputs from the task.
      """
      # Based on the prompt and inputs dictionary, determine the task
      _init(project, domain)
      return (await resources.run_task(name, inputs, project, domain)).to_dict()

• union_mcp/v2/resources.py:4-19 (helper)

  Core helper function that fetches a Flyte task and runs it asynchronously, returning action details. Called by the v2 server handler.

  async def run_task(
      name: str,
      inputs: dict,
      project: str | None = None,
      domain: str | None = None,
      version: str | None = None,
  ) -> flyte.remote.ActionDetails:
      task = flyte.remote.Task.get(
          name=name,
          project=project,
          domain=domain,
          version=version,
          auto_version="latest" if version is None else None,
      )
      run: flyte.remote.Run = flyte.run(task, **inputs)
      return await run.action.details()

• union_mcp/v1/server.py:32-62 (handler)

  v1 MCP tool handler and registration for 'run_task'. Directly fetches and executes task using union remote API.

  @mcp.tool()
  @require_auth
  def run_task(
      name: str,
      inputs: dict,
      project: str,
      domain: str,
      ctx: Context,
  ) -> dict:
      ctx.info(f"Running task {name} in project {project} and domain {domain}")
      """Run a task with natural language.
  
      - Based on the prompt and inputs dictionary, determine the task to run
      - Format the inputs dictionary so that it matches the task function signature
      - Invoke the task
      
      Args:
          project: Project to run the task in.
          domain: Domain to run the task in.
          name: Name of the task to run.
          inputs: A dictionary of inputs to the task.
  
      Returns:
          A dictionary of outputs from the task.
      """
      # Based on the prompt and inputs dictionary, determine the task
      remote = _remote(project, domain)
      task = remote.fetch_task(project=project, domain=domain, name=name)
      execution = remote.execute(task, inputs, project=project, domain=domain)
      return resources.proto_to_json(execution.to_flyte_idl())