implement typed rethrows for higher-order functions

baml_language/crates/baml_builtins2/baml_std/baml/containers.baml

@@ -73,10 +73,18 @@ class Map<K, V> {
   }
 
   function map_keys<U>(self, f: (K) -> U) -> U[] {
-    self.keys().map(f)
+    let result: U[] = []
+    for (let key in self.keys()) {
+      result.push(f(key))
+    }
+    result
   }
 
   function map_values<U>(self, f: (V) -> U) -> U[] {
-    self.values().map(f)
+    let result: U[] = []
+    for (let value in self.values()) {
+      result.push(f(value))
+    }
+    result
   }
 }

baml_language/crates/baml_builtins2/baml_std/testing/registry.baml

@@ -4,13 +4,13 @@
 
 class TestRegistration {
   name string
-  body () -> null
+  body () -> null throws unknown
   runner TestRunner?
 }
 
 class TestSetRegistration {
   name string
-  collector (TestCollector) -> null
+  collector (TestCollector) -> null throws unknown
   runner TestSetRunner?
 }
 
@@ -23,7 +23,7 @@ class TestCollector {
     TestCollector { prefix: prefix, tests: [], testsets: [] }
   }
 
-  function register_test(self, name: string, body: () -> null, runner: TestRunner?) -> null {
+  function register_test(self, name: string, body: () -> null throws unknown, runner: TestRunner?) -> null {
     let full_name = if (self.prefix == "") { name } else { self.prefix + "/" + name }
     let count = 0
     let hash_prefix = full_name + "#"
@@ -35,7 +35,7 @@ class TestCollector {
     null
   }
 
-  function register_test_set(self, name: string, collector: (TestCollector) -> null, runner: TestSetRunner?) -> null {
+  function register_test_set(self, name: string, collector: (TestCollector) -> null throws unknown, runner: TestSetRunner?) -> null {
     let full_name = if (self.prefix == "") { name } else { self.prefix + "/" + name }
     let count = 0
     let hash_prefix = full_name + "#"
@@ -171,7 +171,7 @@ class TestRegistry {
 }
 
 // Run a single test with error handling and optional runner middleware.
-function run_test(body: () -> null, runner: TestRunner?) -> TestReport {
+function run_test(body: () -> null throws unknown, runner: TestRunner?) -> TestReport {
   let base_run = () -> TestReport {
     let start = baml.sys.now_ms()
     let result = {
@@ -187,18 +187,16 @@ function run_test(body: () -> null, runner: TestRunner?) -> TestReport {
       runs: [RunReport { outcome: result.outcome, duration_ms: baml.sys.now_ms() - start }],
     }
   }
-  let effective_run = match (runner) {
-    null => base_run,
-    r: TestRunner => r(base_run)
+  match (runner) {
+    null => base_run(),
+    r: TestRunner => r(base_run)()
   }
-  effective_run()
 }
 
 // Run all children of a testset with optional runner middleware.
 function run_testset(run_children: () -> TestSetReport, runner: TestSetRunner?) -> TestSetReport {
-  let effective_run = match (runner) {
-    null => run_children,
-    r: TestSetRunner => r(run_children)
+  match (runner) {
+    null => run_children(),
+    r: TestSetRunner => r(run_children)()
   }
-  effective_run()
 }

baml_language/crates/baml_builtins2/baml_std/testing/types.baml

@@ -31,7 +31,7 @@ type ChildReport = TestReport | TestSetReport
 
 // Runner type aliases — lambda transformers.
 // TestRunner wraps "run test once" into a new execution lambda.
-type TestRunner = (() -> TestReport) -> () -> TestReport
+type TestRunner = (() -> TestReport throws unknown) -> (() -> TestReport throws unknown)
 
 // TestSetRunner wraps "run all children" into a new execution lambda.
-type TestSetRunner = (() -> TestSetReport) -> () -> TestSetReport
+type TestSetRunner = (() -> TestSetReport throws unknown) -> (() -> TestSetReport throws unknown)

baml_language/crates/baml_compiler2_ast/src/ast.rs

@@ -41,6 +41,7 @@ pub enum TypeExpr {
     /// Named type path: `User`, `baml.http.Request`
     Path {
         segments: Vec<Name>,
+        type_args: Vec<TypeExpr>,
         attrs: Vec<RawAttribute>,
     },
     /// Primitive types
@@ -97,10 +98,11 @@ pub enum TypeExpr {
         value: baml_base::Literal,
         attrs: Vec<RawAttribute>,
     },
-    /// Function type: (params) -> return
+    /// Function type: (params) -> return [throws E]
     Function {
         params: Vec<FunctionTypeParam>,
         ret: Box<TypeExpr>,
+        throws: Option<Box<TypeExpr>>,
         attrs: Vec<RawAttribute>,
     },
     /// The `unknown` keyword type
@@ -445,6 +447,7 @@ pub enum Expr {
     },
     Object {
         type_name: Option<Name>,
+        type_args: Vec<TypeExpr>,
         fields: Vec<(Name, ExprId)>,
         spreads: Vec<SpreadField>,
     },

baml_language/crates/baml_compiler2_ast/src/companions.rs

@@ -100,6 +100,7 @@ fn make_llm_companion(
     let return_type = SpannedTypeExpr {
         expr: TypeExpr::Path {
             segments: return_type_path.iter().map(Name::new).collect(),
+            type_args: vec![],
             attrs: vec![],
         },
         span: parent.span,

baml_language/crates/baml_compiler2_ast/src/disambiguate.rs

@@ -138,11 +138,19 @@ fn validate_type_expr_tree(expr: &TypeExpr, diagnostics: &mut Vec<(String, text_
                 validate_type_expr_tree(v, diagnostics);
             }
         }
-        TypeExpr::Function { params, ret, .. } => {
+        TypeExpr::Function {
+            params,
+            ret,
+            throws,
+            ..
+        } => {
             for p in params {
                 validate_type_expr_tree(&p.ty, diagnostics);
             }
             validate_type_expr_tree(ret, diagnostics);
+            if let Some(throws) = throws {
+                validate_type_expr_tree(throws, diagnostics);
+            }
         }
         _ => {}
     }

baml_language/crates/baml_compiler2_ast/src/lib.rs

@@ -65,6 +65,7 @@ mod tests {
         (Path($name:expr $(, Attr($a:expr))*)) => {
             TypeExpr::Path {
                 segments: vec![baml_base::Name::new($name)],
+                type_args: vec![],
                 attrs: type_expr!(@attrs $(, Attr($a))*),
             }
         };
@@ -153,8 +154,13 @@ mod tests {
             TypeExpr::Rust { attrs } => TypeExpr::Rust {
                 attrs: strip_attrs(attrs),
             },
-            TypeExpr::Path { segments, attrs } => TypeExpr::Path {
+            TypeExpr::Path {
+                segments,
+                type_args,
+                attrs,
+            } => TypeExpr::Path {
                 segments: segments.clone(),
+                type_args: type_args.iter().map(strip_spans).collect(),
                 attrs: strip_attrs(attrs),
             },
             TypeExpr::Optional { inner, attrs } => TypeExpr::Optional {
@@ -178,7 +184,12 @@ mod tests {
                 value: value.clone(),
                 attrs: strip_attrs(attrs),
             },
-            TypeExpr::Function { params, ret, attrs } => TypeExpr::Function {
+            TypeExpr::Function {
+                params,
+                ret,
+                throws,
+                attrs,
+            } => TypeExpr::Function {
                 params: params
                     .iter()
                     .map(|p| crate::ast::FunctionTypeParam {
@@ -187,6 +198,7 @@ mod tests {
                     })
                     .collect(),
                 ret: Box::new(strip_spans(ret)),
+                throws: throws.as_ref().map(|t| Box::new(strip_spans(t))),
                 attrs: strip_attrs(attrs),
             },
             TypeExpr::Media { kind, attrs } => TypeExpr::Media {
@@ -286,6 +298,7 @@ class Response {
                     baml_base::Name::new("errors"),
                     baml_base::Name::new("Io"),
                 ],
+                type_args: vec![],
                 attrs: vec![]
             }
         );
@@ -966,6 +979,7 @@ retry_policy MyRetry {
         let (type_name, fields, _) = match root_expr {
             Expr::Object {
                 type_name,
+                type_args: _,
                 fields,
                 spreads,
             } => (type_name, fields, spreads),

baml_language/crates/baml_compiler2_ast/src/lower_cst.rs

@@ -419,6 +419,7 @@ pub(crate) fn synthesize_llm_builtin_call(
             let counter = alloc(Expr::Literal(Literal::Int(0)));
             alloc(Expr::Object {
                 type_name: Some(Name::new("baml.llm.Client")),
+                type_args: vec![],
                 fields: vec![
                     (Name::new("name"), name_lit),
                     (Name::new("client_type"), ct_variant),
@@ -953,6 +954,7 @@ fn synthesize_init_test_function(
         type_expr: Some(SpannedTypeExpr {
             expr: crate::ast::TypeExpr::Path {
                 segments: vec![Name::new("testing"), Name::new("TestCollector")],
+                type_args: vec![],
                 attrs: vec![],
             },
             span,
@@ -1053,6 +1055,7 @@ fn synthesize_register_call(
                 type_expr: Some(SpannedTypeExpr {
                     expr: crate::ast::TypeExpr::Path {
                         segments: vec![Name::new("testing"), Name::new("TestCollector")],
+                        type_args: vec![],
                         attrs: vec![],
                     },
                     span,
@@ -1229,6 +1232,7 @@ fn synthesize_retry_policy_let(
 
     let root = alloc(Expr::Object {
         type_name: Some(Name::new("RetryPolicy")),
+        type_args: vec![],
         fields,
         spreads: vec![],
     });
@@ -1466,6 +1470,7 @@ fn synthesize_client_let(
     // Client { name, client_type, sub_clients, retry, counter }
     let root = alloc(Expr::Object {
         type_name: Some(Name::new("Client")),
+        type_args: vec![],
         fields: vec![
             (Name::new("name"), name_expr),
             (Name::new("client_type"), client_type_expr),
@@ -1677,6 +1682,7 @@ fn synthesize_client_new_companion(
             if !provider_fields_set.is_empty() {
                 alloc(Expr::Object {
                     type_name: Some(Name::new(type_name)),
+                    type_args: vec![],
                     fields: prov_fields,
                     spreads: vec![],
                 })
@@ -1716,6 +1722,7 @@ fn synthesize_client_new_companion(
 
     let options_expr = alloc(Expr::Object {
         type_name: Some(Name::new("baml.llm.PrimitiveClientOptions")),
+        type_args: vec![],
         fields: options_fields,
         spreads: vec![],
     });
@@ -1727,6 +1734,7 @@ fn synthesize_client_new_companion(
     )));
     let root = alloc(Expr::Object {
         type_name: Some(Name::new("baml.llm.PrimitiveClient")),
+        type_args: vec![],
         fields: vec![
             (Name::new("name"), name_lit),
             (Name::new("provider"), provider_lit),

baml_language/crates/baml_compiler2_ast/src/lower_expr_body.rs

@@ -1088,6 +1088,7 @@ impl LoweringContext {
                                     name,
                                     ty: crate::ast::TypeExpr::Path {
                                         segments: vec![Name::new(&text)],
+                                        type_args: vec![],
                                         attrs: vec![],
                                     },
                                 };
@@ -1990,6 +1991,7 @@ impl LoweringContext {
         let mut spreads = Vec::new();
         let mut position = 0;
         let mut type_name = None;
+        let mut type_args = Vec::new();
 
         // Look for the optional type name (first WORD or path before the brace).
         // The type name may be:
@@ -2021,6 +2023,19 @@ impl LoweringContext {
                     if let Some(name) = last_word {
                         type_name = Some(name);
                     }
+                    type_args = child_node
+                        .children()
+                        .find(|n| n.kind() == SyntaxKind::GENERIC_ARGS)
+                        .into_iter()
+                        .flat_map(|args_node| {
+                            args_node
+                                .children()
+                                .filter(|n| n.kind() == SyntaxKind::TYPE_EXPR)
+                                .filter_map(baml_compiler_syntax::ast::TypeExpr::cast)
+                                .map(|arg| crate::lower_type_expr::lower_type_expr_node(&arg))
+                                .collect::<Vec<_>>()
+                        })
+                        .collect();
                     // After handling the path node, stop scanning for more pre-brace items.
                     break 'outer;
                 }
@@ -2086,6 +2101,7 @@ impl LoweringContext {
         self.alloc_expr(
             Expr::Object {
                 type_name,
+                type_args,
                 fields,
                 spreads,
             },
@@ -2840,6 +2856,7 @@ impl LoweringContext {
             type_expr: Some(SpannedTypeExpr {
                 expr: TypeExpr::Path {
                     segments: vec![Name::new("testing"), Name::new("TestCollector")],
+                    type_args: vec![],
                     attrs: vec![],
                 },
                 span,