rust-codegen.test.ts

import type { ApiSchema } from "../../scripts/codegen/utils.ts";
import type { JSONSchema7 } from "json-schema";
import { describe, expect, it } from "vitest";

import {
    generateApiTypesCode,
    generateRpcCode,
    generateSessionEventsCode,
    isRustCodegenEntrypoint,
} from "../../scripts/codegen/rust.ts";
import { legacyRequestSchema } from "./legacy-parameters-fixture.ts";

describe("Rust codegen entrypoint", () => {
    it("matches Windows paths case-insensitively", () => {
        expect(
            isRustCodegenEntrypoint(
                "C:\\a\\copilot-agent-runtime\\src\\sdk\\scripts\\codegen\\rust.ts",
                "c:\\a\\copilot-agent-runtime\\src\\sdk\\scripts\\codegen\\rust.ts",
                "win32"
            )
        ).toBe(true);
    });
});

describe("Rust API type codegen", () => {
    it.each([true, false])(
        "validates a reference's sibling constant without changing its enum (required: %s)",
        (required) => {
            const code = generateApiTypesCode({
                definitions: {
                    CandidateKind: {
                        type: "string",
                        enum: ["mcp-server", "ai-skill", "unknown"],
                    },
                    Candidate: {
                        type: "object",
                        required: required ? ["kind", "ordinaryKind"] : ["ordinaryKind"],
                        properties: {
                            kind: {
                                $ref: "#/definitions/CandidateKind",
                                const: "ai-skill",
                            },
                            ordinaryKind: { $ref: "#/definitions/CandidateKind" },
                        },
                    },
                },
            } as ApiSchema);

            expect(code).toContain(
                `#[serde(${required ? "" : "default, "}deserialize_with = "Candidate::deserialize_kind")]`
            );
            expect(code).toContain(
                `pub kind: ${required ? "CandidateKind" : "Option<CandidateKind>"},`
            );
            expect(code).toContain("pub ordinary_kind: CandidateKind,");
            expect(code).not.toContain("Candidate::deserialize_ordinary_kind");
            expect(code).toContain(`if value != "ai-skill" {`);
            expect(code).toContain(`&["ai-skill"]`);
            expect(code).toContain(
                "<CandidateKind>::deserialize(serde::de::value::StringDeserializer::<D::Error>::new(value))"
            );
            expect(code.match(/fn deserialize_/g)).toHaveLength(1);
            if (!required) {
                expect(code).toContain("Option::<String>::deserialize(deserializer)?");
                expect(code).toContain("return Ok(None);");
                expect(code).toContain(".map(Some)");
            }
            expect(code).toContain(`#[serde(rename = "mcp-server")]
    McpServer,`);
            expect(code).toContain(`#[serde(rename = "ai-skill")]
    AiSkill,`);
            expect(code).toContain(`#[serde(rename = "unknown")]
    UnknownValue,`);
            expect(code).toContain(`#[default]
    #[serde(other)]
    Unknown,`);
        }
    );

    it("separates adjacent constant helpers before formatting", () => {
        const property: JSONSchema7 = { $ref: "#/definitions/Kind", const: "ai-skill" };
        const code = generateApiTypesCode({
            definitions: {
                Kind: { type: "string", enum: ["ai-skill", "mcp-server"] },
                Candidate: {
                    type: "object",
                    required: ["first", "second"],
                    properties: { first: property, second: { ...property } },
                },
            },
        } as ApiSchema);

        expect(code).toContain("    }\n\n    fn deserialize_second");
    });

    it("distinguishes a protocol-defined unknown value from the forward-compatible fallback", () => {
        const code = generateApiTypesCode({
            definitions: {
                CatalogTrustEligibility: {
                    type: "string",
                    enum: ["default", "expanded", "hidden", "unknown"],
                },
            },
        } as ApiSchema);

        expect(code).toContain(`#[serde(rename = "unknown")]
    UnknownValue,
    /// Unknown variant for forward compatibility.
    #[default]
    #[serde(other)]
    Unknown,`);
    });

    it("publicly re-exports API types moved into the shared session-events schema", () => {
        const code = generateApiTypesCode({
            definitions: {
                PermissionDecision: {
                    type: "object",
                    required: ["source"],
                    properties: {
                        source: {
                            $ref: "session-events.schema.json#/definitions/PermissionDecisionSource",
                        },
                    },
                },
            },
        } as ApiSchema);

        expect(code).toContain("pub use super::session_events::{PermissionDecisionSource};");
        expect(code).toContain("use crate::types::{RequestId, SessionId};");
    });

    it("emits the enqueue result union discriminated by queued", () => {
        const code = generateApiTypesCode({
            definitions: {
                AcceptedEnqueueCommandResult: {
                    type: "object",
                    required: ["queued", "queueId"],
                    properties: {
                        queued: { type: "boolean", const: true },
                        queueId: { type: "string" },
                    },
                },
                UnsupportedEnqueueCommandResult: {
                    type: "object",
                    required: ["queued"],
                    properties: {
                        queued: { type: "boolean", const: false },
                        queueId: { type: ["string", "null"] },
                    },
                },
                EnqueueCommandResult: {
                    anyOf: [
                        { $ref: "#/definitions/AcceptedEnqueueCommandResult" },
                        { $ref: "#/definitions/UnsupportedEnqueueCommandResult" },
                    ],
                    title: "EnqueueCommandResult",
                },
            },
        } as ApiSchema);

        expect(code).toContain("#[serde(untagged)]");
        expect(code).toContain("pub enum EnqueueCommandResult {");
        expect(code).toContain("True(AcceptedEnqueueCommandResult),");
        expect(code).toContain("False(UnsupportedEnqueueCommandResult),");
        expect(code).toContain(
            'deserialize_with = "AcceptedEnqueueCommandResult::deserialize_queued", serialize_with = "AcceptedEnqueueCommandResult::serialize_queued"'
        );
        expect(code).toContain(
            'deserialize_with = "UnsupportedEnqueueCommandResult::deserialize_queued", serialize_with = "UnsupportedEnqueueCommandResult::serialize_queued"'
        );
        expect(code).toContain('serde::de::Error::custom("expected true")');
        expect(code).toContain('serde::ser::Error::custom("expected true")');
        expect(code).toContain('serde::de::Error::custom("expected false")');
        expect(code).toContain('serde::ser::Error::custom("expected false")');
        expect(code).toContain("if !value {");
        expect(code).toContain("if !*value {");
        expect(code).toContain("if value {");
        expect(code).toContain("if *value {");
        expect(code).not.toMatch(/\bvalue != (?:true|false)\b/);
        expect(code).not.toMatch(/\*value != (?:true|false)\b/);
    });

    it.each([
        ["unknown", "unknown_value"],
        ["unknown_value", "unknown"],
        ["not-called", "not_called"],
    ])("rejects colliding wire values %s and %s", (first, second) => {
        expect(() =>
            generateApiTypesCode({
                definitions: {
                    Collision: { type: "string", enum: [first, second] },
                },
            } as ApiSchema)
        ).toThrow("is not unique");
    });

    it("keeps the fallback name reserved for wire values other than exact lowercase unknown", () => {
        expect(() =>
            generateApiTypesCode({
                definitions: {
                    Collision: { type: "string", enum: ["Unknown"] },
                },
            } as ApiSchema)
        ).toThrow('Generated Rust enum variant identifier "Unknown" is not unique');
    });
});

describe("Rust session event codegen", () => {
    function eventSchema(
        data: JSONSchema7,
        definitions: JSONSchema7["definitions"] = {}
    ): JSONSchema7 {
        return {
            definitions: {
                ...definitions,
                SessionEvent: {
                    anyOf: [
                        {
                            type: "object",
                            required: ["type", "data"],
                            properties: {
                                type: { const: "session.search" },
                                data,
                            },
                        },
                    ],
                },
            },
        };
    }

    it.each([
        ["anyOf", false],
        ["anyOf", true],
        ["oneOf", false],
        ["oneOf", true],
    ] as const)("preserves root %s event unions (referenced: %s)", (keyword, referenced) => {
        const payload: JSONSchema7 = {
            ...(referenced ? { title: "SearchPayload" } : {}),
            [keyword]: [
                {
                    type: "object",
                    required: ["kind", "ready"],
                    properties: {
                        kind: { type: "string", const: "status" },
                        ready: { type: "boolean" },
                    },
                },
                {
                    type: "object",
                    required: ["kind", "durationMs"],
                    properties: {
                        kind: { type: "string", const: "startup" },
                        durationMs: { type: "number" },
                    },
                },
            ],
        };
        const code = generateSessionEventsCode(
            eventSchema(
                referenced ? { $ref: "#/definitions/SearchPayload" } : payload,
                referenced ? { SearchPayload: payload } : {}
            )
        );
        const name = referenced ? "SearchPayload" : "SessionSearchData";

        expect(code).toContain("SessionSearch(SessionSearchData),");
        expect(code).toContain(`#[serde(untagged)]
pub enum ${name} {
    Status(${name}Status),
    Startup(${name}Startup),
}`);
        expect(code).toContain("pub ready: bool,");
        expect(code).toContain("pub duration_ms: f64,");
        expect(code).not.toContain("pub struct SessionSearchData {");
        if (referenced) {
            expect(code).toContain("pub type SessionSearchData = SearchPayload;");
        } else {
            expect(code).not.toContain("pub type SessionSearchData = SessionSearchData;");
        }
    });

    it.each(["anyOf", "oneOf"] as const)(
        "preserves unrepresentable root %s event payloads as raw JSON",
        (keyword) => {
            const code = generateSessionEventsCode(
                eventSchema({ [keyword]: [{ type: "string" }, { type: "number" }] })
            );

            expect(code).toContain("pub type SessionSearchData = serde_json::Value;");
            expect(code).toContain("SessionSearch(SessionSearchData),");
            expect(code).not.toContain("pub struct SessionSearchData");
        }
    );

    it("keeps ordinary and empty object event payloads as structs", () => {
        const ordinary = generateSessionEventsCode(
            eventSchema({
                type: "object",
                required: ["message"],
                properties: { message: { type: "string" } },
            })
        );
        const empty = generateSessionEventsCode(eventSchema({ type: "object", properties: {} }));

        expect(ordinary).toContain("pub struct SessionSearchData {\n    pub message: String,\n}");
        expect(empty).toContain("pub struct SessionSearchData {\n}");
        expect(ordinary).not.toContain("pub type SessionSearchData");
        expect(empty).not.toContain("pub type SessionSearchData");
    });

    it.each(["reasonCode", "judgeStatus", "evaluationStage"])(
        "preserves explicit unknown values in nested approval %s enums",
        (property) => {
            const code = generateSessionEventsCode({
                definitions: {
                    SessionEvent: {
                        anyOf: [
                            {
                                type: "object",
                                required: ["type", "data"],
                                properties: {
                                    type: { const: "permission.completed" },
                                    data: {
                                        type: "object",
                                        properties: {
                                            approval: {
                                                type: "object",
                                                properties: {
                                                    [property]: {
                                                        type: "string",
                                                        enum: ["unknown", "inherited"],
                                                    },
                                                },
                                            },
                                        },
                                    },
                                },
                            },
                        ],
                    },
                },
            });

            expect(code).toContain(`#[serde(rename = "unknown")]
    UnknownValue,`);
            expect(code).toContain(`#[serde(rename = "inherited")]
    Inherited,`);
            expect(code).toContain(`/// Unknown variant for forward compatibility.
    #[default]
    #[serde(other)]
    Unknown,`);
        }
    );
});

describe("Rust x-legacy-parameters", () => {
    const render = (additions: 0 | 1 | 2, scope: "server" | "session" = "server") => {
        const { schema } = legacyRequestSchema(additions, scope);
        const types = generateApiTypesCode(schema);
        const rpc = generateRpcCode(schema);
        const block = (code: string, pattern: RegExp) => code.match(pattern)?.[0];
        return {
            types,
            rpc,
            request: block(types, /pub struct SamplePlanRequest \{[\s\S]*?\n\}/),
            constructor: block(types, /pub fn new\([^)]*\) -> Self/),
            plan: block(rpc, /pub async fn plan\([^)]*\)[^{]*/),
            options: block(rpc, /pub async fn plan_with_options\([^)]*\)[^{]*/),
        };
    };

    it("leaves unmarked requests unchanged", () => {
        const original = render(0);
        expect(original.types).not.toContain("SamplePlanOptions");
        expect(original.rpc).not.toContain("plan_with_options");
    });

    it("freezes the published struct and method and adds private-field options", () => {
        const original = render(0);
        const once = render(1);
        expect(once.request).toBe(original.request);
        expect(once.plan).toBe(original.plan);
        expect(once.types).toContain(
            "pub struct SamplePlanOptions {\n    #[serde(flatten)]\n    legacy: SamplePlanRequest,"
        );
        expect(once.types).toMatch(
            /#\[serde\(skip_serializing_if = "Option::is_none"\)\]\n    policy_session_id: Option<String>,/
        );
        expect(once.constructor).toBe(
            "pub fn new(contract: impl Into<String>, source: impl Into<String>) -> Self"
        );
        expect(once.types).toContain(
            "legacy: SamplePlanRequest { contract: contract.into(), source: source.into(), scope: None },"
        );
        expect(once.types).toContain(
            "pub fn scope(mut self, value: impl Into<String>) -> Self {\n        self.legacy.scope = Some(value.into());"
        );
        expect(once.types).toContain(
            "pub fn policy_session_id(mut self, value: impl Into<String>) -> Self {\n        self.policy_session_id = Some(value.into());"
        );
        expect(once.options).toContain("params: SamplePlanOptions");
        const wireCalls = once.rpc.match(/rpc_methods::SAMPLE_PLAN/g) ?? [];
        expect(wireCalls).toHaveLength(2);
        expect(once.types).not.toContain("non_exhaustive");
    });

    it("keeps both entry points unchanged across a second optional addition", () => {
        const once = render(1);
        const twice = render(2);
        expect(twice.request).toBe(once.request);
        expect(twice.constructor).toBe(once.constructor);
        expect(twice.plan).toBe(once.plan);
        expect(twice.options).toBe(once.options);
        expect(twice.types).toContain(
            "pub fn trace_id(mut self, value: impl Into<String>) -> Self {"
        );
        expect(twice.types).toContain("            trace_id: None,");
    });

    it("keeps the session id injected by session-scoped wrappers", () => {
        const once = render(1, "session");
        expect(once.request).toBe(render(0, "session").request);
        expect(once.request).not.toContain("session_id");
        const body = once.rpc.slice(once.rpc.indexOf("pub async fn plan_with_options"));
        expect(body).toMatch(
            /wire_params\["sessionId"\] = serde_json::Value::String\(self\.session\.id\(\)\.to_string\(\)\);/
        );
    });

    it("derives Default for options without required inputs", () => {
        const { schema, params } = legacyRequestSchema(1);
        params.required = [];
        const types = generateApiTypesCode(schema);
        expect(types).toContain("pub fn new() -> Self");
        expect(types).toContain(
            "impl Default for SamplePlanOptions {\n    fn default() -> Self {\n        Self::new()"
        );
        expect(render(1).types).not.toContain("impl Default for SamplePlanOptions");
    });

    it("rejects optional requests", () => {
        const { schema, params } = legacyRequestSchema(1);
        const method = (schema.server as Record<string, Record<string, { params: unknown }>>).sample
            .plan;
        method.params = { anyOf: [{ not: {} }, params] };
        expect(() => generateApiTypesCode(schema)).toThrow(
            "Invalid x-legacy-parameters for sample.plan: optional requests cannot declare legacy parameters"
        );
    });

    it("rejects metadata that omits a required input", () => {
        const { schema, params } = legacyRequestSchema(1);
        (params as Record<string, unknown>)["x-legacy-parameters"] = ["contract", "scope"];
        expect(() => generateApiTypesCode(schema)).toThrow(
            "Invalid x-legacy-parameters for sample.plan: required property source must be a legacy parameter"
        );
    });
});