Expand discriminated unions in action templates

ts/packages/dispatcher/dispatcher/src/translation/actionTemplate.ts

@@ -20,6 +20,7 @@ import {
     ActionParamArray,
     ActionParamObject,
     ActionParamType,
+    resolveTypeReference,
 } from "@typeagent/action-schema";
 import { getActionParamCompletion } from "./requestCompletion.js";
 
@@ -43,31 +44,99 @@ function getDefaultActionTemplate(
     return template;
 }
 
+function resolveObjectType(
+    type: ActionParamType,
+): ActionParamObject | undefined {
+    const resolved = resolveTypeReference(type);
+    return resolved?.type === "object" ? resolved : undefined;
+}
+
+// Discriminated object union: shared single-value kind field.
+function getObjectUnionDiscriminator(
+    types: readonly ActionParamType[],
+):
+    | { fieldName: string; values: string[]; arms: ActionParamObject[] }
+    | undefined {
+    if (types.length < 2) {
+        return undefined;
+    }
+    const arms: ActionParamObject[] = [];
+    for (const t of types) {
+        const obj = resolveObjectType(t);
+        if (obj === undefined) {
+            return undefined;
+        }
+        arms.push(obj);
+    }
+
+    const firstFields = Object.keys(arms[0].fields);
+    for (const fieldName of firstFields) {
+        const values: string[] = [];
+        let ok = true;
+        for (const arm of arms) {
+            const field = arm.fields[fieldName];
+            if (field === undefined || field.optional) {
+                ok = false;
+                break;
+            }
+            const ft = resolveTypeReference(field.type) ?? field.type;
+            if (ft.type !== "string-union" || ft.typeEnum.length !== 1) {
+                ok = false;
+                break;
+            }
+            values.push(ft.typeEnum[0]);
+        }
+        if (!ok) {
+            continue;
+        }
+        // Kind values must be unique per arm.
+        if (new Set(values).size !== values.length) {
+            continue;
+        }
+        return { fieldName, values, arms };
+    }
+    return undefined;
+}
+
 function toTemplateTypeObject(
     type: ActionParamObject,
     visited: ReadonlySet<string>,
+    data: unknown,
 ) {
     const templateType: TemplateFieldObject = {
         type: "object",
         fields: {},
     };
 
+    const dataObj =
+        data !== null && typeof data === "object" && !Array.isArray(data)
+            ? (data as Record<string, unknown>)
+            : undefined;
+
     for (const [key, field] of Object.entries(type.fields)) {
-        const type = toTemplateType(field.type, visited);
-        if (type === undefined) {
+        const fieldData = dataObj !== undefined ? dataObj[key] : undefined;
+        const fieldType = toTemplateType(field.type, visited, fieldData);
+        if (fieldType === undefined) {
             // Skip undefined fields.
             continue;
         }
-        templateType.fields[key] = { optional: field.optional, type };
+        templateType.fields[key] = {
+            optional: field.optional,
+            type: fieldType,
+        };
     }
     return templateType;
 }
 
 function toTemplateTypeArray(
     type: ActionParamArray,
     visited: ReadonlySet<string>,
+    data: unknown,
 ) {
-    const elementType = toTemplateType(type.elementType, visited);
+    // First array element shapes nested templates.
+    const elementData =
+        Array.isArray(data) && data.length > 0 ? data[0] : undefined;
+    const elementType = toTemplateType(type.elementType, visited, elementData);
     if (elementType === undefined) {
         // Skip undefined fields.
         return undefined;
@@ -79,14 +148,92 @@ function toTemplateTypeArray(
     return templateType;
 }
 
+// Expand discriminated unions; else first-arm fallback.
+function toTemplateTypeUnion(
+    types: readonly ActionParamType[],
+    visited: ReadonlySet<string>,
+    data: unknown,
+): TemplateType | undefined {
+    const disc = getObjectUnionDiscriminator(types);
+    if (disc !== undefined) {
+        const { fieldName, values, arms } = disc;
+        let selectedIndex = 0;
+        if (
+            data !== null &&
+            typeof data === "object" &&
+            !Array.isArray(data) &&
+            fieldName in (data as object)
+        ) {
+            const current = (data as Record<string, unknown>)[fieldName];
+            if (typeof current === "string") {
+                const idx = values.indexOf(current);
+                if (idx >= 0) {
+                    selectedIndex = idx;
+                }
+            }
+        }
+        const selectedArm = arms[selectedIndex];
+        const template = toTemplateTypeObject(selectedArm, visited, data);
+        // Full kind enum; arm fields from data.
+        template.fields[fieldName] = {
+            optional: false,
+            type: {
+                type: "string-union",
+                typeEnum: values,
+                discriminator: values[selectedIndex],
+            },
+        };
+        return template;
+    }
+
+    // Prefer arm that matches current data shape.
+    if (data !== undefined) {
+        for (const t of types) {
+            const resolved = resolveTypeReference(t) ?? t;
+            try {
+                if (
+                    resolved.type === "object" &&
+                    data !== null &&
+                    typeof data === "object" &&
+                    !Array.isArray(data)
+                ) {
+                    return toTemplateType(t, visited, data);
+                }
+                if (resolved.type === "array" && Array.isArray(data)) {
+                    return toTemplateType(t, visited, data);
+                }
+                if (
+                    (resolved.type === "string" ||
+                        resolved.type === "number" ||
+                        resolved.type === "boolean" ||
+                        resolved.type === "string-union") &&
+                    typeof data === resolved.type
+                ) {
+                    return toTemplateType(t, visited, data);
+                }
+                if (
+                    resolved.type === "string-union" &&
+                    typeof data === "string"
+                ) {
+                    return toTemplateType(t, visited, data);
+                }
+            } catch {
+                // try next arm
+            }
+        }
+    }
+    // Historical fallback: first arm.
+    return toTemplateType(types[0], visited, data);
+}
+
 function toTemplateType(
     type: ActionParamType,
     visited: ReadonlySet<string> = new Set<string>(),
+    data: unknown = undefined,
 ): TemplateType | undefined {
     switch (type.type) {
         case "type-union":
-            // TODO: smarter about type unions.
-            return toTemplateType(type.types[0], visited);
+            return toTemplateTypeUnion(type.types, visited, data);
         case "type-reference":
             if (type.definition === undefined) {
                 throw new Error(`Unresolved type reference: ${type.name}`);
@@ -98,11 +245,12 @@ function toTemplateType(
             return toTemplateType(
                 type.definition.type,
                 new Set([...visited, type.name]),
+                data,
             );
         case "object":
-            return toTemplateTypeObject(type, visited);
+            return toTemplateTypeObject(type, visited, data);
         case "array":
-            return toTemplateTypeArray(type, visited);
+            return toTemplateTypeArray(type, visited, data);
         case "string-union":
             return type as TemplateFieldStringUnion;
         case "string":
@@ -153,7 +301,12 @@ function toTemplate(
 
     const actionParametersType = actionSchema.type.fields.parameters?.type;
     if (actionParametersType) {
-        const type = toTemplateType(actionParametersType);
+        // Pass data so union arms resolve correctly.
+        const type = toTemplateType(
+            actionParametersType,
+            new Set(),
+            action.parameters,
+        );
         if (type !== undefined) {
             template.fields.parameters = {
                 // ActionParam types are compatible with TemplateFields

ts/packages/dispatcher/dispatcher/test/actionTemplate.spec.ts

@@ -92,3 +92,134 @@ describe("getSystemTemplateSchema toTemplateType", () => {
         expect(Object.keys(parameters.fields)).toEqual(["keep"]);
     });
 });
+
+describe("getSystemTemplateSchema discriminated type-union", () => {
+    function musicTargetUnion(): ActionParamType {
+        const track = {
+            type: "object",
+            fields: {
+                kind: {
+                    type: { type: "string-union", typeEnum: ["track"] },
+                },
+                trackName: { type: { type: "string" } },
+                artists: {
+                    optional: true,
+                    type: {
+                        type: "array",
+                        elementType: { type: "string" },
+                    },
+                },
+            },
+        };
+        const artist = {
+            type: "object",
+            fields: {
+                kind: {
+                    type: { type: "string-union", typeEnum: ["artist"] },
+                },
+                artistName: { type: { type: "string" } },
+            },
+        };
+        const album = {
+            type: "object",
+            fields: {
+                kind: {
+                    type: { type: "string-union", typeEnum: ["album"] },
+                },
+                albumName: { type: { type: "string" } },
+            },
+        };
+        const playlist = {
+            type: "object",
+            fields: {
+                kind: {
+                    type: { type: "string-union", typeEnum: ["playlist"] },
+                },
+                playlistName: { type: { type: "string" } },
+            },
+        };
+        return {
+            type: "object",
+            fields: {
+                target: {
+                    type: {
+                        type: "type-union",
+                        types: [track, artist, album, playlist],
+                    },
+                },
+            },
+        } as unknown as ActionParamType;
+    }
+
+    it("exposes full kind enum for MusicTarget-style discriminated unions", async () => {
+        const template = await getSystemTemplateSchema(
+            "action",
+            {
+                schemaName,
+                actionName,
+                parameters: {
+                    target: { kind: "track", trackName: "Bohemian Rhapsody" },
+                },
+            },
+            makeContext(musicTargetUnion()),
+        );
+
+        const parameters = template.fields.parameters?.type as any;
+        const target = parameters.fields.target.type;
+        expect(target.type).toBe("object");
+        expect(target.fields.kind.type).toEqual({
+            type: "string-union",
+            typeEnum: ["track", "artist", "album", "playlist"],
+            discriminator: "track",
+        });
+        expect(target.fields.trackName.type).toEqual({ type: "string" });
+        expect(target.fields.artistName).toBeUndefined();
+    });
+
+    it("selects album arm fields when data.kind is album", async () => {
+        const template = await getSystemTemplateSchema(
+            "action",
+            {
+                schemaName,
+                actionName,
+                parameters: {
+                    target: {
+                        kind: "album",
+                        albumName: "A Night at the Opera",
+                    },
+                },
+            },
+            makeContext(musicTargetUnion()),
+        );
+
+        const target = (template.fields.parameters?.type as any).fields.target
+            .type;
+        expect(target.fields.kind.type.typeEnum).toEqual([
+            "track",
+            "artist",
+            "album",
+            "playlist",
+        ]);
+        expect(target.fields.kind.type.discriminator).toBe("album");
+        expect(target.fields.albumName.type).toEqual({ type: "string" });
+        expect(target.fields.trackName).toBeUndefined();
+        expect(target.fields.playlistName).toBeUndefined();
+    });
+
+    it("keeps non-union object parameters unchanged", async () => {
+        const parametersType: ActionParamType = {
+            type: "object",
+            fields: {
+                query: { type: { type: "string" } },
+            },
+        } as unknown as ActionParamType;
+
+        const template = await getSystemTemplateSchema(
+            "action",
+            { ...data, parameters: { query: "hello" } },
+            makeContext(parametersType),
+        );
+        const parameters = template.fields.parameters?.type as any;
+        expect(parameters.fields.query.type).toEqual({ type: "string" });
+    });
+});