Refactor paramsSchema to function-based pattern with z in context

packages/cli-core/src/features/DeployCommand/DeployCommand.ts

@@ -19,6 +19,7 @@ export interface IDeployCommandParams {
     deploymentLogs?: boolean;
     build?: boolean;
     preview?: boolean;
+    allowLocalStateFiles?: boolean;
 }
 
 export interface IDeploySingleAppParams {
@@ -82,6 +83,12 @@ export class DeployCommand implements CliCommand.Interface<IDeployCommandParams>
                     description: "Print deployment logs (automatically enabled in CI environments)",
                     type: "boolean",
                     default: false
+                },
+                {
+                    name: "allow-local-state-files",
+                    description:
+                        "Allow using local Pulumi state files with production environment deployment (not recommended).",
+                    type: "boolean"
                 }
             ],
             handler: async (params: IDeployCommandParams) => {

packages/project/src/abstractions/features/DeployApp.ts

@@ -9,6 +9,7 @@ interface IDeployAppParams {
     preview?: boolean;
     debug?: boolean;
     dataMigrationLogStreaming?: boolean;
+    allowLocalStateFiles?: boolean;
     output?: (pulumiProcess: IPulumiProcess) => Promise<void>;
 }
 

packages/project/src/abstractions/models/IProjectConfigModel.ts

@@ -1,12 +1,12 @@
 import { type IHydratedProjectConfig } from "~/abstractions/models/IProjectConfigDto.js";
 import { type ExtensionInstanceModel } from "~/defineExtension/models/index.js";
-import { type z } from "zod";
+import { type ParamsSchemaDefinition } from "~/defineExtension/types.js";
 import { type ExtensionComponent } from "~/defineExtension/index.js";
 
 export interface IProjectConfigModel {
     config: IHydratedProjectConfig;
 
-    extensionsByType<TParamsSchema extends z.ZodTypeAny>(
+    extensionsByType<TParamsSchema extends ParamsSchemaDefinition | undefined>(
         type: string | ExtensionComponent<TParamsSchema>
     ): ExtensionInstanceModel<TParamsSchema>[];
 }

packages/project/src/abstractions/services/IsRemotePulumiBackendService.ts

@@ -0,0 +1,13 @@
+import { createAbstraction } from "~/abstractions/createAbstraction.js";
+
+export interface IIsRemotePulumiBackendService {
+    execute(): boolean;
+}
+
+export const IsRemotePulumiBackendService = createAbstraction<IIsRemotePulumiBackendService>(
+    "IsRemotePulumiBackendService"
+);
+
+export namespace IsRemotePulumiBackendService {
+    export type Interface = IIsRemotePulumiBackendService;
+}

packages/project/src/abstractions/services/index.ts

@@ -9,6 +9,7 @@ export { GetProjectIdService } from "./GetProjectIdService.js";
 export { GetProjectService } from "./GetProjectService.js";
 export { GetProjectVersionService } from "./GetProjectVersionService.js";
 export { GetPulumiService } from "./GetPulumiService.js";
+export { IsRemotePulumiBackendService } from "./IsRemotePulumiBackendService.js";
 export { ListAppLambdaFunctionsService } from "./ListAppLambdaFunctionsService.js";
 export { ListDeployedEnvironmentsService } from "./ListDeployedEnvironmentsService.js";
 export { ListPackagesInAppWorkspaceService } from "./ListPackagesInAppWorkspaceService.js";

packages/project/src/createProjectSdkContainer.ts

@@ -63,6 +63,7 @@ import {
     getPulumiService,
     getPulumiVersionService,
     getYarnVersionService,
+    isRemotePulumiBackendService,
     listAppLambdaFunctionsService,
     listDeployedEnvironmentsService,
     listPackagesInAppWorkspaceService,
@@ -151,6 +152,7 @@ export const createProjectSdkContainer = async (
     container.register(getPulumiService).inSingletonScope();
     container.register(getPulumiVersionService).inSingletonScope();
     container.register(getYarnVersionService).inSingletonScope();
+    container.register(isRemotePulumiBackendService).inSingletonScope();
     container.register(listAppLambdaFunctionsService).inSingletonScope();
     container.register(listDeployedEnvironmentsService).inSingletonScope();
     container.register(listPackagesInAppWorkspaceService).inSingletonScope();

packages/project/src/defineExtension/createExtensionDefinition.ts

@@ -1,8 +1,8 @@
 import { ExtensionDefinitionModel } from "./models/index.js";
 import type { DefineExtensionParams } from "./types.js";
-import type { z } from "zod";
+import type { ParamsSchemaDefinition } from "./types.js";
 
-export function createExtensionDefinition<TParamsSchema extends z.ZodTypeAny>(
+export function createExtensionDefinition<TParamsSchema extends ParamsSchemaDefinition | undefined>(
     extensionParams: DefineExtensionParams<TParamsSchema>
 ) {
     const { type, description, multiple, build, validate, tags, paramsSchema } = extensionParams;

packages/project/src/defineExtension/createExtensionReactComponent.tsx

@@ -1,7 +1,6 @@
 import React, { useMemo } from "react";
 import { Property, useIdGenerator } from "@webiny/react-properties";
-import { type DefineExtensionParams } from "./types.js";
-import { type z } from "zod";
+import { type DefineExtensionParams, type ParamsSchemaDefinition, type ParamsSchemaInfer } from "./types.js";
 
 const KeyValues = (props: Record<string, any>) => {
     const getId = useIdGenerator("");
@@ -10,14 +9,15 @@ const KeyValues = (props: Record<string, any>) => {
     });
 };
 
-type ExtensionReactComponentProps<TParamsSchema extends z.ZodTypeAny> = z.infer<TParamsSchema> & {
+type ExtensionReactComponentProps<TParamsSchema extends ParamsSchemaDefinition | undefined> = 
+    (TParamsSchema extends ParamsSchemaDefinition ? ParamsSchemaInfer<TParamsSchema> : {}) & {
     remove?: boolean;
     before?: string;
     after?: string;
     name?: string;
 };
 
-export function createExtensionReactComponent<TParamsSchema extends z.ZodTypeAny>(
+export function createExtensionReactComponent<TParamsSchema extends ParamsSchemaDefinition | undefined>(
     extensionParams: DefineExtensionParams<TParamsSchema>
 ) {
     const ExtensionReactComponent: React.FC<

packages/project/src/defineExtension/defineApiExtension.ts

@@ -18,7 +18,7 @@ export const defineApiExtension = (params: DefineApiExtensionParams) =>
         tags: { runtimeContext: "app-build", appName: "api" },
         description: params.description,
         multiple: true,
-        paramsSchema: ({ project }) => {
+        paramsSchema: ({ project, z }) => {
             if (!params.abstraction) {
                 return z.object({
                     src: z.string(),

packages/project/src/defineExtension/defineExtension.ts

@@ -1,16 +1,15 @@
-import { type DefineExtensionParams } from "~/defineExtension/types.js";
+import { type DefineExtensionParams, type ParamsSchemaDefinition } from "~/defineExtension/types.js";
 import { createExtensionDefinition } from "./createExtensionDefinition.js";
 import { createExtensionReactComponent } from "./createExtensionReactComponent.js";
-import { type z } from "zod";
 
-export type ExtensionComponent<TParamsSchema extends z.ZodTypeAny> = ReturnType<
+export type ExtensionComponent<TParamsSchema extends ParamsSchemaDefinition | undefined> = ReturnType<
     typeof createExtensionReactComponent<TParamsSchema>
 > & {
     def: ReturnType<typeof createExtensionDefinition<TParamsSchema>>;
     getDefinition: () => ReturnType<typeof createExtensionDefinition<TParamsSchema>>;
 };
 
-export function defineExtension<TParamsSchema extends z.ZodTypeAny>(
+export function defineExtension<TParamsSchema extends ParamsSchemaDefinition | undefined>(
     extensionParams: DefineExtensionParams<TParamsSchema>
 ) {
     const definition = createExtensionDefinition<TParamsSchema>(extensionParams);

packages/project/src/defineExtension/models/ExtensionDefinitionModel.ts

@@ -1,29 +1,37 @@
-import { type ExtensionTags } from "~/defineExtension/types.js";
-import { type z } from "zod";
-import { type ExtensionInstanceModelContext } from "~/defineExtension/index.js";
+import { type ExtensionTags, type ParamsSchemaDefinition, type ParamsSchemaInfer, type ParamsSchemaFunction } from "~/defineExtension/types.js";
 
-export interface ExtensionDefinitionModelParams<TParamsSchema extends z.ZodTypeAny> {
+export interface ExtensionDefinitionModelParams<TParamsSchema extends ParamsSchemaDefinition | undefined> {
     type: string;
     tags?: ExtensionTags;
     description: string;
     array?: boolean;
-    paramsSchema?: TParamsSchema | ((ctx: ExtensionInstanceModelContext) => TParamsSchema);
+    paramsSchema?: ParamsSchemaFunction<TParamsSchema>;
 
-    build?(params: TParamsSchema, ctx: ExtensionInstanceModelContext): Promise<void> | void;
+    build?(
+        params: TParamsSchema extends ParamsSchemaDefinition ? ParamsSchemaInfer<TParamsSchema> : any,
+        ctx: any
+    ): Promise<void> | void;
 
-    validate?(params: TParamsSchema): Promise<void> | void;
+    validate?(
+        params: TParamsSchema extends ParamsSchemaDefinition ? ParamsSchemaInfer<TParamsSchema> : any
+    ): Promise<void> | void;
 }
 
-export class ExtensionDefinitionModel<TParamsSchema extends z.ZodTypeAny> {
+export class ExtensionDefinitionModel<TParamsSchema extends ParamsSchemaDefinition | undefined> {
     type: string;
     description: string;
     tags: ExtensionTags;
     multiple?: boolean;
-    paramsSchema?: TParamsSchema | ((ctx: ExtensionInstanceModelContext) => TParamsSchema);
+    paramsSchema?: ParamsSchemaFunction<TParamsSchema>;
 
-    build?(params: TParamsSchema, ctx: ExtensionInstanceModelContext): Promise<void> | void;
+    build?(
+        params: TParamsSchema extends ParamsSchemaDefinition ? ParamsSchemaInfer<TParamsSchema> : any,
+        ctx: any
+    ): Promise<void> | void;
 
-    validate?(params: TParamsSchema): Promise<void> | void;
+    validate?(
+        params: TParamsSchema extends ParamsSchemaDefinition ? ParamsSchemaInfer<TParamsSchema> : any
+    ): Promise<void> | void;
 
     constructor(params: ExtensionDefinitionModelParams<TParamsSchema>) {
         this.type = params.type;

packages/project/src/defineExtension/models/ExtensionInstanceModel.ts

@@ -1,18 +1,19 @@
 import { type ExtensionDefinitionModel } from "./ExtensionDefinitionModel.js";
-import { type z } from "zod";
+import { z } from "zod";
 import { type ProjectModel } from "~/models/index.js";
 import { ProjectError } from "~/ProjectError.js";
+import { type ParamsSchemaDefinition, type ParamsSchemaInfer } from "~/defineExtension/types.js";
 
 export interface ExtensionInstanceModelContext {
     [key: string]: any;
 
     project: ProjectModel;
 }
 
-export class ExtensionInstanceModel<TParamsSchema extends z.ZodTypeAny> {
+export class ExtensionInstanceModel<TParamsSchema extends ParamsSchemaDefinition | undefined> {
     constructor(
         public definition: ExtensionDefinitionModel<TParamsSchema>,
-        public params: z.infer<TParamsSchema>,
+        public params: TParamsSchema extends ParamsSchemaDefinition ? ParamsSchemaInfer<TParamsSchema> : any,
         public context: ExtensionInstanceModelContext
     ) {}
 
@@ -29,10 +30,19 @@ export class ExtensionInstanceModel<TParamsSchema extends z.ZodTypeAny> {
             return;
         }
 
-        const paramsSchema =
-            typeof this.definition.paramsSchema === "function"
-                ? this.definition.paramsSchema(this.context)
-                : this.definition.paramsSchema;
+        // Call paramsSchema with context that includes z
+        const contextWithZ = { ...this.context, z };
+        const result = this.definition.paramsSchema(contextWithZ);
+
+        // Check if the result is already a Zod object or a plain object
+        let paramsSchema: z.ZodObject<any>;
+        if (result && typeof result === 'object' && 'safeParse' in result) {
+            // Result is already a ZodObject (e.g., z.object({...}))
+            paramsSchema = result as z.ZodObject<any>;
+        } else {
+            // Result is a plain object, wrap it
+            paramsSchema = z.object(result as any);
+        }
 
         const validationResult = await paramsSchema.safeParseAsync(this.params);
         if (!validationResult.success) {

packages/project/src/defineExtension/types.ts

@@ -8,15 +8,27 @@ export type ExtensionTags = {
     runtimeContext?: "app-build" | "project" | "cli" | "pulumi";
 };
 
-export interface DefineExtensionParams<TParamsSchema extends z.ZodTypeAny> {
+export type ParamsSchemaDefinition = Record<string, z.ZodTypeAny>;
+
+export type ParamsSchemaInfer<T extends ParamsSchemaDefinition | undefined> =
+    T extends ParamsSchemaDefinition ? z.infer<z.ZodObject<T> & { [k: string]: any }> : never;
+
+// Context with z included for the new pattern
+export type ParamsSchemaContext = ExtensionInstanceModelContext & { z: typeof z };
+
+// Type for the new unified pattern: ({ z, project, ... }) => ({...}) or z.object({...})
+export type ParamsSchemaFunction<TParamsSchema extends ParamsSchemaDefinition | undefined> = 
+    (ctx: ParamsSchemaContext) => TParamsSchema | z.ZodObject<any>;
+
+export interface DefineExtensionParams<TParamsSchema extends ParamsSchemaDefinition | undefined> {
     type: string;
     tags: ExtensionTags;
     description?: string;
     multiple?: boolean;
-    paramsSchema?: TParamsSchema | ((ctx: ExtensionInstanceModelContext) => TParamsSchema);
+    paramsSchema?: ParamsSchemaFunction<TParamsSchema>;
     build?: (
-        params: z.infer<TParamsSchema>,
+        params: TParamsSchema extends ParamsSchemaDefinition ? ParamsSchemaInfer<TParamsSchema> : any,
         ctx: ExtensionInstanceModelContext
     ) => Promise<void> | void;
-    validate?: (params: z.infer<TParamsSchema>) => Promise<void> | void;
+    validate?: (params: TParamsSchema extends ParamsSchemaDefinition ? ParamsSchemaInfer<TParamsSchema> : any) => Promise<void> | void;
 }

packages/project/src/extensions/DatabaseSetup.ts

@@ -1,11 +1,10 @@
-import { z } from "zod";
 import { defineExtension } from "../defineExtension/index.js";
 
 export const DatabaseSetup = defineExtension({
     type: "Project/DatabaseSetup",
     tags: { runtimeContext: "project" },
     description: "Define the database setup configuration (ddb, ddb+es, or ddb+os).",
-    paramsSchema: z.object({
+    paramsSchema: ({ z }) => ({
         setupName: z
             .enum(["ddb", "ddb+es", "ddb+os"])
             .describe(

packages/project/src/extensions/EnvVar.ts

@@ -1,12 +1,11 @@
-import { z } from "zod";
 import { defineExtension } from "~/defineExtension/index.js";
 
 export const EnvVar = defineExtension({
     type: "Project/EnvVar",
     tags: { runtimeContext: "project" },
     description: "Set an environment variable in the project context.",
     multiple: true,
-    paramsSchema: z.object({
+    paramsSchema: ({ z }) => ({
         // TODO: enable using `name` instead of `varName` for better consistency.
         varName: z.string().describe("The environment variable name."),
         value: z.string().describe("The environment variable value.")

packages/project/src/extensions/ExtensionDefinitions.ts

@@ -1,12 +1,11 @@
 import { defineExtension } from "~/defineExtension/index.js";
-import { z } from "zod";
 
 export const ExtensionDefinitions = defineExtension({
     type: "Project/ExtensionDefinitions",
     tags: { runtimeContext: "project" },
     description: "Register additional extension definitions.",
     multiple: true,
-    paramsSchema: z.object({
+    paramsSchema: ({ z }) => ({
         src: z.string()
     })
 });

packages/project/src/extensions/Project.tsx

@@ -9,7 +9,8 @@ import {
     ApiBeforeWatch,
     BeforeBuild,
     BeforeDeploy,
-    BeforeWatch
+    BeforeWatch,
+    CoreBeforeDeploy
 } from "~/extensions/index.js";
 
 const p = createPathResolver(import.meta.dirname, "Project");
@@ -28,6 +29,7 @@ export const Project = () => {
             <BeforeBuild src={p("WcpSetEnvVarsBeforeBuild.js")} />
             <BeforeWatch src={p("WcpSetEnvVarsBeforeWatch.js")} />
             <BeforeDeploy src={p("EnsureTelemetryEnabledForOss.js")} />
+            <CoreBeforeDeploy src={p("ValidateProductionPulumiState.js")} />
         </>
     );
 };

packages/project/src/extensions/Project/ValidateProductionPulumiState.ts

@@ -0,0 +1,49 @@
+import {
+    CoreBeforeDeploy,
+    IsRemotePulumiBackendService,
+    ProjectSdkParamsService
+} from "~/abstractions/index.js";
+import { GracefulError } from "@webiny/project";
+
+class ValidateProductionPulumiStateImpl implements CoreBeforeDeploy.Interface {
+    constructor(
+        private isRemotePulumiBackendService: IsRemotePulumiBackendService.Interface,
+        private projectSdkParamsService: ProjectSdkParamsService.Interface
+    ) {}
+
+    async execute(params: CoreBeforeDeploy.Params) {
+        const sdkParams = this.projectSdkParamsService.get();
+        const { env } = sdkParams;
+        const { allowLocalStateFiles } = params;
+
+        const prodEnvs = ["prod", "production"];
+        const isProdEnv = prodEnvs.includes(env);
+
+        if (!isProdEnv) {
+            return;
+        }
+
+        if (this.isRemotePulumiBackendService.execute()) {
+            return;
+        }
+
+        if (allowLocalStateFiles) {
+            return;
+        }
+
+        const error = new Error("Cannot deploy to production with local state files.");
+
+        const message = [
+            "Use the %s flag to continue with local Pulumi state files,",
+            "or configure a remote backend for production deployments.",
+            "Learn more: https://webiny.link/state-files-production."
+        ].join(" ");
+
+        throw GracefulError.from(error, message, "--allow-local-state-files");
+    }
+}
+
+export const ValidateProductionPulumiState = CoreBeforeDeploy.createImplementation({
+    implementation: ValidateProductionPulumiStateImpl,
+    dependencies: [IsRemotePulumiBackendService, ProjectSdkParamsService]
+});

packages/project/src/extensions/ProjectDecorator.ts

@@ -7,7 +7,7 @@ export const ProjectDecorator = defineExtension({
     tags: { runtimeContext: "project" },
     multiple: true,
     description: "Decorate an existing implementation with additional functionality.",
-    paramsSchema: ({ project }) => {
+    paramsSchema: ({ project, z }) => {
         return z.object({
             src: zodPathToFile(project)
         });