From a627d96a865ee9a4a6c1088159edd522e5a99f2f Mon Sep 17 00:00:00 2001
From: Jian Hui <v-jileong@microsoft.com>
Date: Mon, 21 Sep 2026 14:28:25 +0800
Subject: [PATCH] fix: link cls references when generating commands from
 typespec

Typespec operations are parsed from json, so their cls references arrive
without `implement` set, while the swagger builder links them while it
builds the schema. `generate_outputs()` runs before `CMDCommand.link()`,
so a response body that is a bare cls reference - e.g. the body carrying 202
of the ARM LRO put of Microsoft.Compute galleries - hit an AssertionError in
`_build_output_type_schema`.
---
 .../swagger/model/schema/cmd_builder.py       |  6 +-
 .../test_typespec_command_generator.py        | 76 +++++++++++++++++++
 2 files changed, 81 insertions(+), 1 deletion(-)
 create mode 100644 src/aaz_dev/swagger/tests/controller_tests/test_typespec_command_generator.py

diff --git a/src/aaz_dev/swagger/model/schema/cmd_builder.py b/src/aaz_dev/swagger/model/schema/cmd_builder.py
index e7ff744a..df87c749 100644
--- a/src/aaz_dev/swagger/model/schema/cmd_builder.py
+++ b/src/aaz_dev/swagger/model/schema/cmd_builder.py
@@ -650,7 +650,11 @@ def apply_cls_definitions(self, *cmd_ops):
                 cmd_op.register_cls(cls_register_map=schema_cls_register_map)
             
             for name, cls_register in schema_cls_register_map.items():
-                if cls_register.get('implement', None):
+                if implement := cls_register.get('implement', None):
+                    # typespec operations are parsed from json, so their cls references arrive
+                    # unlinked (swagger links them while building the schema).
+                    for refer in cls_register['refers']:
+                        refer.implement = implement
                     continue
 
                 new_schema = None
diff --git a/src/aaz_dev/swagger/tests/controller_tests/test_typespec_command_generator.py b/src/aaz_dev/swagger/tests/controller_tests/test_typespec_command_generator.py
new file mode 100644
index 00000000..f90b6840
--- /dev/null
+++ b/src/aaz_dev/swagger/tests/controller_tests/test_typespec_command_generator.py
@@ -0,0 +1,76 @@
+import unittest
+
+from command.model.configuration import CMDResource
+from swagger.controller.command_generator import TypespecCommandGenerator
+from utils.base64 import b64encode_str
+
+PATH = "/subscriptions/{subscriptionId}/resourceGroups/{resourceGroupName}/providers/Microsoft.Mock/galleries/{galleryName}"
+VERSION = "2026-03-03"
+
+
+def _put_path_item():
+    # Shape emitted by typespec-aaz for an ARM LRO put whose 202 repeats the resource body
+    # (e.g. Microsoft.Compute galleries): the first response owns the definition and the
+    # following one is a bare cls reference.
+    resource_schema = {
+        "type": "object",
+        "cls": "Gallery_read",
+        "props": [
+            {"type": "string", "name": "id", "readOnly": True},
+            {"type": "string", "name": "location", "required": True},
+        ],
+    }
+    return {
+        "put": {
+            "operationId": "Galleries_CreateOrUpdate",
+            "create": {
+                "operationId": "Galleries_CreateOrUpdate",
+                "http": {
+                    "path": PATH,
+                    "request": {
+                        "method": "put",
+                        "path": {"params": [
+                            {"type": "string", "name": "galleryName", "required": True},
+                            {"type": "string", "name": "resourceGroupName", "required": True},
+                            {"type": "string", "name": "subscriptionId", "required": True},
+                        ]},
+                        "query": {"params": [
+                            {"type": "string", "name": "api-version", "required": True},
+                        ]},
+                        "body": {"json": {"schema": {
+                            "type": "object",
+                            "name": "resource",
+                            "required": True,
+                            "clientFlatten": True,
+                            "props": [{"type": "string", "name": "location", "required": True}],
+                        }}},
+                    },
+                    "responses": [
+                        {"statusCode": [200, 201], "body": {"json": {"schema": resource_schema}}},
+                        {"statusCode": [202], "body": {"json": {"schema": {"type": "@Gallery_read"}}}},
+                        {"isError": True, "body": {"json": {"schema": {
+                            "readOnly": True, "type": "@MgmtErrorFormat"}}}},
+                    ],
+                },
+            },
+        },
+    }
+
+
+class TypespecCommandGeneratorTestCase(unittest.TestCase):
+
+    def test_output_from_cls_reference_response(self):
+        generator = TypespecCommandGenerator()
+        generator.load_resources([{"path": PATH, "pathItem": _put_path_item()}])
+        resource = CMDResource({
+            "id": "/subscriptions/{}/resourcegroups/{}/providers/microsoft.mock/galleries/{}",
+            "version": VERSION,
+            "swagger": "mgmt-plane/mock/ResourceProviders/Microsoft.Mock"
+                       f"/Paths/{b64encode_str(PATH)}/V/{b64encode_str(VERSION)}",
+        })
+
+        command_group = generator.create_draft_command_group(
+            resource, instance_var="$Instance", methods=("put",))
+
+        create_command = command_group.commands[0]
+        self.assertEqual([output.type for output in create_command.outputs], ["object"])