chore: refactoring things to work a little cleaner

Signed-off-by: blam <ben@blam.sh>

Signed-off-by: blam <ben@blam.sh>

plugins/scaffolder-node/src/actions/createTemplateAction.test.ts

@@ -0,0 +1,125 @@
+/*
+ * Copyright 2025 The Backstage Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import { createTemplateAction } from './createTemplateAction';
+import { z } from 'zod';
+
+describe('createTemplateAction', () => {
+  it('should allow creating with jsonschema and use the old deprecated types', () => {
+    const action = createTemplateAction<{ repoUrl: string }, { test: string }>({
+      id: 'test',
+      schema: {
+        input: {
+          type: 'object',
+          required: ['repoUrl'],
+          properties: {
+            repoUrl: { type: 'string' },
+          },
+        },
+        output: {
+          type: 'object',
+          required: ['test'],
+          properties: {
+            test: { type: 'string' },
+          },
+        },
+      },
+      handler: async ctx => {
+        // @ts-expect-error - repoUrl is string
+        const a: number = ctx.input.repoUrl;
+
+        const b: string = ctx.input.repoUrl;
+        expect(b).toBeDefined();
+
+        const stream = ctx.logStream;
+        expect(stream).toBeDefined();
+
+        ctx.output('test', 'value');
+
+        // @ts-expect-error - not valid output type
+        ctx.output('test', 4);
+
+        // @ts-expect-error - not valid output name
+        ctx.output('test2', 'value');
+      },
+    });
+  });
+
+  it('should allow creating with zod and use the old deprecated types', () => {
+    const action = createTemplateAction({
+      id: 'test',
+      schema: {
+        input: z.object({
+          repoUrl: z.string(),
+        }),
+        output: z.object({
+          test: z.string(),
+        }),
+      },
+      handler: async ctx => {
+        // @ts-expect-error - repoUrl is string
+        const a: number = ctx.input.repoUrl;
+
+        const b: string = ctx.input.repoUrl;
+        expect(b).toBeDefined();
+
+        const stream = ctx.logStream;
+        expect(stream).toBeDefined();
+
+        ctx.output('test', 'value');
+
+        // @ts-expect-error - not valid output type
+        ctx.output('test', 4);
+
+        // @ts-expect-error - not valid output name
+        ctx.output('test2', 'value');
+      },
+    });
+  });
+
+  it('should allow creating with new first class zod support', () => {
+    const action = createTemplateAction({
+      id: 'test',
+      schema: {
+        input: {
+          repoUrl: d => d.string(),
+        },
+        output: {
+          test: d => d.string(),
+        },
+      },
+      handler: async ctx => {
+        // @ts-expect-error - repoUrl is string
+        const a: number = ctx.input.repoUrl;
+
+        const b: string = ctx.input.repoUrl;
+        expect(b).toBeDefined();
+
+        // @ts-expect-error - logStream is not available
+        const stream = ctx.logStream;
+
+        expect(stream).toBeDefined();
+
+        ctx.output('test', 'value');
+
+        // @ts-expect-error - not valid output type
+        ctx.output('test', 4);
+
+        // @ts-expect-error - not valid output name
+        ctx.output('test2', 'value');
+      },
+    });
+  });
+});

plugins/scaffolder-node/src/actions/createTemplateAction.ts

@@ -17,8 +17,8 @@
 import { ActionContext, TemplateAction } from './types';
 import { z } from 'zod';
 import { Schema } from 'jsonschema';
-import zodToJsonSchema from 'zod-to-json-schema';
-import { JsonObject } from '@backstage/types';
+import { Expand, JsonObject } from '@backstage/types';
+import { parseSchemas } from './util';
 
 /** @public */
 export type TemplateExample = {
@@ -30,8 +30,14 @@ export type TemplateExample = {
 export type TemplateActionOptions<
   TActionInput extends JsonObject = {},
   TActionOutput extends JsonObject = {},
-  TInputSchema extends Schema | z.ZodType = {},
-  TOutputSchema extends Schema | z.ZodType = {},
+  TInputSchema extends
+    | Schema
+    | z.ZodType
+    | { [key in string]: (zImpl: typeof z) => z.ZodType } = Schema,
+  TOutputSchema extends
+    | Schema
+    | z.ZodType
+    | { [key in string]: (zImpl: typeof z) => z.ZodType } = Schema,
 > = {
   id: string;
   description?: string;
@@ -41,7 +47,9 @@ export type TemplateActionOptions<
     input?: TInputSchema;
     output?: TOutputSchema;
   };
-  handler: (ctx: ActionContext<TActionInput, TActionOutput>) => Promise<void>;
+  handler: (
+    ctx: ActionContext<TActionInput, TActionOutput, TInputSchema>,
+  ) => Promise<void>;
 };
 
 /**
@@ -49,17 +57,86 @@ export type TemplateActionOptions<
  * Will convert zod schemas to json schemas for use throughout the system.
  * @public
  */
-export const createTemplateAction = <
+export function createTemplateAction<
+  TInputSchema extends { [key in string]: (zImpl: typeof z) => z.ZodType },
+  TOutputSchema extends { [key in string]: (zImpl: typeof z) => z.ZodType },
+>(
+  action: TemplateActionOptions<
+    {
+      [key in keyof TInputSchema]: z.infer<ReturnType<TInputSchema[key]>>;
+    },
+    {
+      [key in keyof TOutputSchema]: z.infer<ReturnType<TOutputSchema[key]>>;
+    },
+    TInputSchema,
+    TOutputSchema
+  >,
+): TemplateAction<
+  Expand<{
+    [key in keyof TInputSchema]: z.output<ReturnType<TInputSchema[key]>>;
+  }>,
+  Expand<{
+    [key in keyof TOutputSchema]: z.output<ReturnType<TOutputSchema[key]>>;
+  }>,
+  TInputSchema
+>;
+/**
+ * @deprecated migrate to using the new built in zod schema definitions for schemas
+ */
+export function createTemplateAction<
   TInputParams extends JsonObject = JsonObject,
   TOutputParams extends JsonObject = JsonObject,
-  TInputSchema extends Schema | z.ZodType = {},
-  TOutputSchema extends Schema | z.ZodType = {},
+  TInputSchema extends Schema = Schema,
+  TOutputSchema extends Schema = Schema,
+  TActionInput extends JsonObject = TInputParams,
+  TActionOutput extends JsonObject = TOutputParams,
+>(
+  action: TemplateActionOptions<
+    TActionInput,
+    TActionOutput,
+    TInputSchema,
+    TOutputSchema
+  >,
+): TemplateAction<TActionInput, TActionOutput, TInputSchema>;
+/**
+ * @deprecated migrate to using the new built in zod schema definitions for schemas
+ */
+export function createTemplateAction<
+  TInputParams extends JsonObject = JsonObject,
+  TOutputParams extends JsonObject = JsonObject,
+  TInputSchema extends z.ZodType = z.ZodType,
+  TOutputSchema extends z.ZodType = z.ZodType,
+  TActionInput extends JsonObject = z.infer<TInputSchema>,
+  TActionOutput extends JsonObject = z.infer<TOutputSchema>,
+>(
+  action: TemplateActionOptions<
+    TActionInput,
+    TActionOutput,
+    TInputSchema,
+    TOutputSchema
+  >,
+): TemplateAction<TActionInput, TActionOutput, TInputSchema>;
+export function createTemplateAction<
+  TInputParams extends JsonObject = JsonObject,
+  TOutputParams extends JsonObject = JsonObject,
+  TInputSchema extends
+    | Schema
+    | z.ZodType
+    | { [key in string]: (zImpl: typeof z) => z.ZodType } = {},
+  TOutputSchema extends
+    | Schema
+    | z.ZodType
+    | { [key in string]: (zImpl: typeof z) => z.ZodType } = {},
   TActionInput extends JsonObject = TInputSchema extends z.ZodType<
     any,
     any,
     infer IReturn
   >
     ? IReturn
+    : TInputSchema extends { [key in string]: (zImpl: typeof z) => z.ZodType }
+    ? Expand<{
+        [key in keyof TInputSchema]: z.infer<ReturnType<TInputSchema[key]>>;
+      }>
     : TInputParams,
   TActionOutput extends JsonObject = TOutputSchema extends z.ZodType<
     any,
@@ -67,6 +144,10 @@ export const createTemplateAction = <
     infer IReturn
   >
     ? IReturn
+    : TOutputSchema extends { [key in string]: (zImpl: typeof z) => z.ZodType }
+    ? Expand<{
+        [key in keyof TOutputSchema]: z.infer<ReturnType<TOutputSchema[key]>>;
+      }>
     : TOutputParams,
 >(
   action: TemplateActionOptions<
@@ -75,16 +156,8 @@ export const createTemplateAction = <
     TInputSchema,
     TOutputSchema
   >,
-): TemplateAction<TActionInput, TActionOutput> => {
-  const inputSchema =
-    action.schema?.input && 'safeParseAsync' in action.schema.input
-      ? zodToJsonSchema(action.schema.input)
-      : action.schema?.input;
-
-  const outputSchema =
-    action.schema?.output && 'safeParseAsync' in action.schema.output
-      ? zodToJsonSchema(action.schema.output)
-      : action.schema?.output;
+): TemplateAction<TActionInput, TActionOutput, TInputSchema> {
+  const { inputSchema, outputSchema } = parseSchemas(action);
 
   return {
     ...action,
@@ -94,4 +167,4 @@ export const createTemplateAction = <
       output: outputSchema as TOutputSchema,
     },
   };
-};
+}

plugins/scaffolder-node/src/actions/types.ts

@@ -21,8 +21,11 @@ import { TaskSecrets } from '../tasks';
 import { TemplateInfo } from '@backstage/plugin-scaffolder-common';
 import { UserEntity } from '@backstage/catalog-model';
 import { Schema } from 'jsonschema';
-import { BackstageCredentials } from '@backstage/backend-plugin-api';
-
+import {
+  BackstageCredentials,
+  LoggerService,
+} from '@backstage/backend-plugin-api';
+import { z } from 'zod';
 /**
  * ActionContext is passed into scaffolder actions.
  * @public
@@ -30,77 +33,148 @@ import { BackstageCredentials } from '@backstage/backend-plugin-api';
 export type ActionContext<
   TActionInput extends JsonObject,
   TActionOutput extends JsonObject = JsonObject,
-> = {
-  // TODO(blam): move this to LoggerService
-  logger: Logger;
-  /** @deprecated - use `ctx.logger` instead */
-  logStream: Writable;
-  secrets?: TaskSecrets;
-  workspacePath: string;
-  input: TActionInput;
-  checkpoint<T extends JsonValue | void>(opts: {
-    key: string;
-    fn: () => Promise<T> | T;
-  }): Promise<T>;
-  output(
-    name: keyof TActionOutput,
-    value: TActionOutput[keyof TActionOutput],
-  ): void;
+  TInputSchema extends
+    | { [key in string]: (zImpl: typeof z) => z.ZodType }
+    | Schema = Schema,
+> = TInputSchema extends { [key in string]: (zImpl: typeof z) => z.ZodType }
+  ? {
+      logger: LoggerService;
+      secrets?: TaskSecrets;
+      workspacePath: string;
+      input: TActionInput;
+      checkpoint<U extends JsonValue>(
+        key: string,
+        fn: () => Promise<U>,
+      ): Promise<U>;
+      output(
+        ...params: {
+          [K in keyof TActionOutput]: [name: K, value: TActionOutput[K]];
+        }[keyof TActionOutput]
+      ): void;
+      /**
+       * Creates a temporary directory for use by the action, which is then cleaned up automatically.
+       */
+      createTemporaryDirectory(): Promise<string>;
 
-  /**
-   * Creates a temporary directory for use by the action, which is then cleaned up automatically.
-   */
-  createTemporaryDirectory(): Promise<string>;
+      /**
+       * Get the credentials for the current request
+       */
+      getInitiatorCredentials(): Promise<BackstageCredentials>;
 
-  /**
-   * Get the credentials for the current request
-   */
-  getInitiatorCredentials(): Promise<BackstageCredentials>;
+      /**
+       * Task information
+       */
+      task: {
+        id: string;
+      };
 
-  /**
-   * Task information
-   */
-  task: {
-    id: string;
-  };
+      templateInfo?: TemplateInfo;
 
-  templateInfo?: TemplateInfo;
+      /**
+       * Whether this action invocation is a dry-run or not.
+       * This will only ever be true if the actions as marked as supporting dry-runs.
+       */
+      isDryRun?: boolean;
 
-  /**
-   * Whether this action invocation is a dry-run or not.
-   * This will only ever be true if the actions as marked as supporting dry-runs.
-   */
-  isDryRun?: boolean;
+      /**
+       * The user which triggered the action.
+       */
+      user?: {
+        /**
+         * The decorated entity from the Catalog
+         */
+        entity?: UserEntity;
+        /**
+         * An entity ref for the author of the task
+         */
+        ref?: string;
+      };
 
-  /**
-   * The user which triggered the action.
-   */
-  user?: {
-    /**
-     * The decorated entity from the Catalog
-     */
-    entity?: UserEntity;
-    /**
-     * An entity ref for the author of the task
-     */
-    ref?: string;
-  };
+      /**
+       * Implement the signal to make your custom step abortable https://developer.mozilla.org/en-US/docs/Web/API/AbortController/signal
+       */
+      signal?: AbortSignal;
 
-  /**
-   * Implement the signal to make your custom step abortable https://developer.mozilla.org/en-US/docs/Web/API/AbortController/signal
-   */
-  signal?: AbortSignal;
+      /**
+       * Optional value of each invocation
+       */
+      each?: JsonObject;
+    }
+  : /** @deprecated **/
+    {
+      // TODO(blam): move this to LoggerService
+      logger: Logger;
+      /** @deprecated - use `ctx.logger` instead */
+      logStream: Writable;
+      secrets?: TaskSecrets;
+      workspacePath: string;
+      input: TActionInput;
+      checkpoint<T extends JsonValue | void>(opts: {
+        key: string;
+        fn: () => Promise<T> | T;
+      }): Promise<T>;
+      output(
+        name: keyof TActionOutput,
+        value: TActionOutput[keyof TActionOutput],
+      ): void;
 
-  /**
-   * Optional value of each invocation
-   */
-  each?: JsonObject;
-};
+      /**
+       * Creates a temporary directory for use by the action, which is then cleaned up automatically.
+       */
+      createTemporaryDirectory(): Promise<string>;
+
+      /**
+       * Get the credentials for the current request
+       */
+      getInitiatorCredentials(): Promise<BackstageCredentials>;
+
+      /**
+       * Task information
+       */
+      task: {
+        id: string;
+      };
+
+      templateInfo?: TemplateInfo;
+
+      /**
+       * Whether this action invocation is a dry-run or not.
+       * This will only ever be true if the actions as marked as supporting dry-runs.
+       */
+      isDryRun?: boolean;
+
+      /**
+       * The user which triggered the action.
+       */
+      user?: {
+        /**
+         * The decorated entity from the Catalog
+         */
+        entity?: UserEntity;
+        /**
+         * An entity ref for the author of the task
+         */
+        ref?: string;
+      };
+
+      /**
+       * Implement the signal to make your custom step abortable https://developer.mozilla.org/en-US/docs/Web/API/AbortController/signal
+       */
+      signal?: AbortSignal;
+
+      /**
+       * Optional value of each invocation
+       */
+      each?: JsonObject;
+    };
 
 /** @public */
 export type TemplateAction<
   TActionInput extends JsonObject = JsonObject,
   TActionOutput extends JsonObject = JsonObject,
+  TInputSchema extends
+    | { [key in string]: (zImpl: typeof z) => z.ZodType }
+    | Schema = Schema,
 > = {
   id: string;
   description?: string;
@@ -110,5 +184,7 @@ export type TemplateAction<
     input?: Schema;
     output?: Schema;
   };
-  handler: (ctx: ActionContext<TActionInput, TActionOutput>) => Promise<void>;
+  handler: (
+    ctx: ActionContext<TActionInput, TActionOutput, TInputSchema>,
+  ) => Promise<void>;
 };

plugins/scaffolder-node/src/actions/util.ts

@@ -18,6 +18,9 @@ import { InputError } from '@backstage/errors';
 import { isChildPath } from '@backstage/backend-plugin-api';
 import { join as joinPath, normalize as normalizePath } from 'path';
 import { ScmIntegrationRegistry } from '@backstage/integration';
+import { TemplateActionOptions } from './createTemplateAction';
+import zodToJsonSchema from 'zod-to-json-schema';
+import { z } from 'zod';
 
 /**
  * @public
@@ -124,3 +127,52 @@ function checkRequiredParams(repoUrl: URL, ...params: string[]) {
     }
   }
 }
+
+const isZodSchema = (schema: unknown): schema is z.ZodType => {
+  return typeof schema === 'object' && !!schema && 'safeParseAsync' in schema;
+};
+
+const isNativeZodSchema = (
+  schema: unknown,
+): schema is { [key in string]: (zImpl: typeof z) => z.ZodType } => {
+  return (
+    typeof schema === 'object' &&
+    !!schema &&
+    Object.values(schema).every(v => typeof v === 'function')
+  );
+};
+
+export const parseSchemas = (action: TemplateActionOptions<any, any, any>) => {
+  if (!action.schema) {
+    return { inputSchema: undefined, outputSchema: undefined };
+  }
+
+  if (isZodSchema(action.schema.input) && isZodSchema(action.schema.output)) {
+    return {
+      inputSchema: zodToJsonSchema(action.schema.input),
+      outputSchema: zodToJsonSchema(action.schema.output),
+    };
+  }
+
+  if (
+    isNativeZodSchema(action.schema.input) &&
+    isNativeZodSchema(action.schema.output)
+  ) {
+    const input = z.object(
+      Object.fromEntries(
+        Object.entries(action.schema.input).map(([k, v]) => [k, v(z)]),
+      ),
+    );
+    const output = z.object(
+      Object.fromEntries(
+        Object.entries(action.schema.output).map(([k, v]) => [k, v(z)]),
+      ),
+    );
+    return {
+      inputSchema: zodToJsonSchema(input),
+      outputSchema: zodToJsonSchema(output),
+    };
+  }
+
+  throw new Error('Invalid schema provided');
+};