11import asyncio
22import json
3+ from typing import Any , Callable , overload
34import uuid
45
56from uipath ._cli .middlewares import MiddlewareResult
67
78from ._utils ._config import McpConfig
89
910
10- async def mcp_init_middleware_async (entrypoint : str ) -> MiddlewareResult :
11+ async def mcp_init_middleware_async (
12+ entrypoint : str ,
13+ options : dict [str , Any ] | None = None ,
14+ write_config : Callable [[Any ], str ] | None = None ,
15+ ) -> MiddlewareResult :
1116 """Middleware to check for mcp.json and create uipath.json with schemas"""
17+ options = options or {}
18+
1219 config = McpConfig ()
1320 if not config .exists :
1421 return MiddlewareResult (
@@ -39,10 +46,13 @@ async def mcp_init_middleware_async(entrypoint: str) -> MiddlewareResult:
3946 "entryPoints" : entrypoints
4047 }
4148
42- config_path = "uipath.json"
43-
44- with open (config_path , "w" ) as f :
45- json .dump (uipath_data , f , indent = 4 )
49+ if write_config :
50+ config_path = write_config (uipath_data )
51+ else :
52+ # Save the uipath.json file
53+ config_path = "uipath.json"
54+ with open (config_path , "w" ) as f :
55+ json .dump (uipath_data , f , indent = 2 )
4656
4757 return MiddlewareResult (
4858 should_continue = False ,
@@ -56,7 +66,23 @@ async def mcp_init_middleware_async(entrypoint: str) -> MiddlewareResult:
5666 should_include_stacktrace = True ,
5767 )
5868
69+ @overload
70+ def mcp_init_middleware (entrypoint : str ) -> MiddlewareResult : ...
71+
72+
73+ @overload
74+ def mcp_init_middleware (
75+ entrypoint : str ,
76+ options : dict [str , Any ],
77+ write_config : Callable [[Any ], str ],
78+ ) -> MiddlewareResult : ...
79+
80+
5981
60- def mcp_init_middleware (entrypoint : str ) -> MiddlewareResult :
82+ def mcp_init_middleware (
83+ entrypoint : str ,
84+ options : dict [str , Any ] | None = None ,
85+ write_config : Callable [[Any ], str ] | None = None ,
86+ ) -> MiddlewareResult :
6187 """Middleware to check for mcp.json and create uipath.json with schemas"""
62- return asyncio .run (mcp_init_middleware_async (entrypoint ))
88+ return asyncio .run (mcp_init_middleware_async (entrypoint , options , write_config ))
0 commit comments