fix: register packaged extensions before model_ast import

hugr-core/src/envelope/reader.rs

@@ -205,7 +205,6 @@ impl<R: BufRead> EnvelopeReader<R> {
         } else {
             ExtensionRegistry::new([])
         };
-        self.register_packaged(&packaged_extensions);
 
         self.import_package(&model_package, packaged_extensions)
             .map_err(Into::into)
@@ -248,6 +247,7 @@ impl<R: BufRead> EnvelopeReader<R> {
         packaged_extensions: ExtensionRegistry,
     ) -> Result<Package, crate::import::ImportError> {
         self.description.set_n_modules(package.modules.len());
+        self.register_packaged(&packaged_extensions);
 
         let modules = package
             .modules
@@ -322,7 +322,7 @@ pub(crate) enum ModelBinaryReadError {
 mod test {
     use super::*;
 
-    use crate::extension::ExtensionRegistry;
+    use crate::extension::{ExtensionId, ExtensionRegistry};
 
     use crate::envelope::header::EnvelopeHeader;
     use cool_asserts::assert_matches;
@@ -466,4 +466,41 @@ mod test {
         handle_error(&mut desc, &error);
         assert!(desc.used_extensions_resolved.is_none());
     }
+
+    #[test]
+    fn test_decode_model_ast_with_packaged_extensions() {
+        let mut simple_package = crate::builder::test::simple_package();
+        // Create a simple extension to package
+        let ext_id = ExtensionId::new("test.packaged.extension").unwrap();
+        let extension = Extension::new(ext_id.clone(), crate::extension::Version::new(1, 0, 0));
+
+        simple_package
+            .extensions
+            .register(std::sync::Arc::new(extension))
+            .unwrap();
+
+        let header = EnvelopeHeader {
+            format: EnvelopeFormat::ModelTextWithExtensions,
+            ..Default::default()
+        };
+
+        let mut cursor = Cursor::new(Vec::new());
+        simple_package.store(&mut cursor, header.config()).unwrap();
+        cursor.set_position(0);
+
+        let registry = ExtensionRegistry::new([]);
+        let mut reader = EnvelopeReader::new(cursor, &registry).unwrap();
+
+        // Before decoding, the packaged extension should not be registered
+        assert!(!reader.registry.contains(&ext_id));
+
+        let result = reader.decode_model_ast();
+
+        // After decoding, the packaged extension should be registered
+        assert!(result.is_ok());
+        assert!(reader.registry.contains(&ext_id));
+
+        let package = result.unwrap();
+        assert!(package.extensions.contains(&ext_id));
+    }
 }