feat(go): expand generated SDK completeness

baml_language/crates/baml_cli/Cargo.toml

@@ -18,6 +18,7 @@ workspace = true
 
 [dependencies]
 baml_builtins2 = { workspace = true }
+baml_codegen_types = { workspace = true }
 baml_db = { workspace = true }
 baml_exec = { workspace = true }
 baml_fmt = { workspace = true }

baml_language/crates/baml_cli/src/generate.rs

@@ -7,6 +7,7 @@ use std::{
 };
 
 use anyhow::{Context, Result, anyhow};
+use baml_codegen_types::{GeneratedOutputFile, write_generated_output};
 use baml_db::{
     FileId, Span,
     baml_compiler_diagnostics::{Diagnostic, DiagnosticId, DiagnosticPhase, Severity, render},
@@ -42,6 +43,7 @@ struct GeneratorDef {
     /// Required for Go so generated packages can import the SDK root and one
     /// another. Other generators leave this unset.
     sdk_import_path: Option<String>,
+    max_typed_union_arity: usize,
 }
 
 impl GenerateArgs {
@@ -247,15 +249,19 @@ impl GenerateArgs {
                         .map(|(path, content)| (path, content.into_bytes()))
                         .collect()
                 }
-                OutputType::Go => sdkgen_go::to_source_code_with_bytecode(
+                OutputType::Go => sdkgen_go::try_to_source_code_with_bytecode_and_options(
                     &pool,
                     &baml_bytecode,
-                    generator.naming_convention,
-                    generator
-                        .sdk_import_path
-                        .as_deref()
-                        .expect("validated Go generator must have sdk_import_path"),
+                    &sdkgen_go::GoGenOptions {
+                        naming_convention: generator.naming_convention,
+                        sdk_import_path: generator
+                            .sdk_import_path
+                            .as_deref()
+                            .expect("validated Go generator must have sdk_import_path"),
+                        max_typed_union_arity: generator.max_typed_union_arity,
+                    },
                 )
+                .map_err(|error| anyhow!("Go SDK generation failed: {error}"))?
                 .into_iter()
                 .map(|(path, content)| (path, content.into_bytes()))
                 .collect(),
@@ -293,24 +299,17 @@ impl GenerateArgs {
                 OutputType::CSharp => unreachable!("C# generation commits atomically above"),
             };
 
-            std::fs::create_dir_all(&output_dir).with_context(|| {
+            let output = generated
+                .into_iter()
+                .map(|(path, contents)| GeneratedOutputFile::new(path, contents))
+                .collect();
+            let report = write_generated_output(&output_dir, output).with_context(|| {
                 format!(
-                    "Failed to create output directory: {}",
+                    "Failed to install generated output in {}",
                     output_dir.display()
                 )
             })?;
-            let output_dir = std::fs::canonicalize(&output_dir).unwrap_or(output_dir);
-
-            let mut count = 0;
-            for (rel_path, content) in &generated {
-                let dest = output_dir.join(rel_path);
-                if let Some(parent) = dest.parent() {
-                    std::fs::create_dir_all(parent)?;
-                }
-                std::fs::write(&dest, content)
-                    .with_context(|| format!("Failed to write {}", dest.display()))?;
-                count += 1;
-            }
+            let count = report.written_files.len();
 
             // Persistent status line in the scrollback — one per
             // generator block. Matches cargo's `   Compiling foo
@@ -399,6 +398,51 @@ fn discover_generators(root: &Path) -> (Vec<GeneratorDef>, Vec<Diagnostic>) {
         } else {
             None
         };
+        let max_typed_union_arity = if matches!(output_type, Some(OutputType::Go)) {
+            match generator.max_typed_union_arity.as_ref() {
+                None => sdkgen_go::DEFAULT_MAX_TYPED_UNION_ARITY,
+                Some(value) if *value.get_ref() >= 0 => *value.get_ref() as usize,
+                Some(value) => {
+                    diags.push(
+                        Diagnostic::error(
+                            DiagnosticId::InvalidGeneratorPropertyValue,
+                            format!(
+                                "Go generator `{name}` requires `max_typed_union_arity` to be zero or greater"
+                            ),
+                        )
+                        .with_primary(
+                            Span {
+                                file_id: manifest_file_id(),
+                                range: to_text_range(value.span()),
+                            },
+                            "negative union threshold",
+                        )
+                        .with_phase(DiagnosticPhase::Validation),
+                    );
+                    sdkgen_go::DEFAULT_MAX_TYPED_UNION_ARITY
+                }
+            }
+        } else {
+            if let Some(value) = generator.max_typed_union_arity.as_ref() {
+                diags.push(
+                    Diagnostic::error(
+                        DiagnosticId::InvalidGeneratorPropertyValue,
+                        format!(
+                            "generator `{name}` sets Go-only property `max_typed_union_arity` on a non-Go target"
+                        ),
+                    )
+                    .with_primary(
+                        Span {
+                            file_id: manifest_file_id(),
+                            range: to_text_range(value.span()),
+                        },
+                        "remove this Go-only property",
+                    )
+                    .with_phase(DiagnosticPhase::Validation),
+                );
+            }
+            sdkgen_go::DEFAULT_MAX_TYPED_UNION_ARITY
+        };
 
         // `output_dir` is resolved relative to the project root and defaults
         // to "..", with the target-owned generated directory appended.
@@ -447,6 +491,7 @@ fn discover_generators(root: &Path) -> (Vec<GeneratorDef>, Vec<Diagnostic>) {
             output_dir,
             naming_convention,
             sdk_import_path,
+            max_typed_union_arity,
         });
     }
 
@@ -585,7 +630,24 @@ fn to_text_range(span: std::ops::Range<usize>) -> TextRange {
 
 #[cfg(test)]
 mod tests {
-    use super::is_valid_go_import_path;
+    use std::fs;
+
+    use super::{Diagnostic, GeneratorDef, discover_generators, is_valid_go_import_path};
+
+    fn go_manifest(threshold: Option<i64>) -> String {
+        let threshold = threshold
+            .map(|value| format!("max_typed_union_arity = {value}\n"))
+            .unwrap_or_default();
+        format!(
+            "[package]\nname = \"test\"\n\n[generator.go]\noutput_type = \"go\"\nnaming_convention = \"language\"\nsdk_import_path = \"example.com/test/baml_sdk\"\n{threshold}"
+        )
+    }
+
+    fn discover_with_manifest(content: &str) -> (Vec<GeneratorDef>, Vec<Diagnostic>) {
+        let directory = tempfile::tempdir().unwrap();
+        fs::write(directory.path().join("baml.toml"), content).unwrap();
+        discover_generators(directory.path())
+    }
 
     #[test]
     fn go_import_paths_reject_relative_empty_and_platform_specific_segments() {
@@ -605,4 +667,36 @@ mod tests {
         }
         assert!(is_valid_go_import_path("example.com/project/baml_sdk"));
     }
+
+    #[test]
+    fn go_union_threshold_defaults_to_three_and_accepts_zero() {
+        let (defaults, default_diags) = discover_with_manifest(&go_manifest(None));
+        assert!(default_diags.is_empty(), "{default_diags:?}");
+        assert_eq!(defaults[0].max_typed_union_arity, 3);
+
+        let (disabled, disabled_diags) = discover_with_manifest(&go_manifest(Some(0)));
+        assert!(disabled_diags.is_empty(), "{disabled_diags:?}");
+        assert_eq!(disabled[0].max_typed_union_arity, 0);
+    }
+
+    #[test]
+    fn negative_go_union_threshold_is_rejected() {
+        let (_, diagnostics) = discover_with_manifest(&go_manifest(Some(-1)));
+        assert_eq!(diagnostics.len(), 1, "{diagnostics:?}");
+        assert!(
+            format!("{diagnostics:?}").contains("zero or greater"),
+            "{diagnostics:?}"
+        );
+    }
+
+    #[test]
+    fn go_union_threshold_on_non_go_generator_is_rejected() {
+        let manifest = "[package]\nname = \"test\"\n\n[generator.ts]\noutput_type = \"typescript/node\"\nnaming_convention = \"language\"\nmax_typed_union_arity = 3\n";
+        let (_, diagnostics) = discover_with_manifest(manifest);
+        assert_eq!(diagnostics.len(), 1, "{diagnostics:?}");
+        assert!(
+            format!("{diagnostics:?}").contains("Go-only property"),
+            "{diagnostics:?}"
+        );
+    }
 }

baml_language/crates/baml_cli/src/init_command.rs

@@ -203,6 +203,7 @@ name = "{name}"
 # output_dir = "."
 # naming_convention = "language"
 # sdk_import_path = "example.com/my-project/baml_sdk"
+# max_typed_union_arity = 3 # larger unions are generated as `any`
 "#,
     )
 }

baml_language/crates/baml_cli/src/manifest.rs

@@ -110,6 +110,10 @@ pub(crate) struct GeneratorManifest {
     /// generated subpackages must import one another by module path.
     pub sdk_import_path: Option<Spanned<String>>,
 
+    /// Maximum non-null union arity represented as a closed generated Go
+    /// union. Larger unions use `any`. Go-only; defaults to 3.
+    pub max_typed_union_arity: Option<Spanned<i64>>,
+
     #[serde(flatten)]
     pub unknown: IndexMap<String, toml::Value>,
 }
@@ -275,4 +279,19 @@ mod tests {
                 .contains("cannot be empty")
         );
     }
+
+    #[test]
+    fn parses_zero_go_union_threshold_with_a_value_span() {
+        let manifest = parse(
+            "[generator.go]\noutput_type = \"go\"\nnaming_convention = \"language\"\nmax_typed_union_arity = 0\n",
+        )
+        .unwrap();
+        let threshold = manifest.generator["go"]
+            .get_ref()
+            .max_typed_union_arity
+            .as_ref()
+            .unwrap();
+        assert_eq!(*threshold.get_ref(), 0);
+        assert!(threshold.span().end > threshold.span().start);
+    }
 }

baml_language/crates/baml_cli/tests/exit_code_e2e.rs

@@ -44,6 +44,10 @@ fn run_baml_cli(built: &Path, dir: &Path, args: &[&str]) -> Output {
     cmd.output().expect("spawn baml-cli")
 }
 
+fn gofmt_is_available() -> bool {
+    Command::new("gofmt").arg("-h").output().is_ok()
+}
+
 /// Create a minimal project structure with the given source code.
 fn create_project(dir: &Path, source: &str) {
     std::fs::write(
@@ -276,6 +280,9 @@ fn generate_valid_project_returns_zero_exit_code() {
 
 #[test]
 fn generate_go_writes_sdk_through_cli() {
+    if !gofmt_is_available() {
+        return;
+    }
     let built = &common::baml_cli();
     let tmp = tempfile::tempdir().unwrap();
     create_project_with_go_generator(
@@ -302,6 +309,84 @@ fn generate_go_writes_sdk_through_cli() {
         std::fs::read_to_string(tmp.path().join("baml_sdk/internal/bootstrap/bootstrap.go"))
             .expect("Go bootstrap should be generated");
     assert!(bootstrap.contains("func Ensure() error"));
+    assert_eq!(
+        std::fs::read_to_string(tmp.path().join("baml_sdk/.gitignore")).unwrap(),
+        baml_codegen_types::GENERATED_GITIGNORE
+    );
+}
+
+#[test]
+fn generate_go_first_run_preserves_preexisting_user_files() {
+    if !gofmt_is_available() {
+        return;
+    }
+    let built = &common::baml_cli();
+    let tmp = tempfile::tempdir().unwrap();
+    create_project_with_go_generator(
+        tmp.path(),
+        "function echo(value: string) -> string { value }\n",
+    );
+    let sdk = tmp.path().join("baml_sdk");
+    std::fs::create_dir(&sdk).unwrap();
+    std::fs::write(sdk.join("user-notes.txt"), "keep me").unwrap();
+
+    let output = run_baml_cli(built, tmp.path(), &["generate", "--from", "."]);
+
+    assert!(
+        output.status.success(),
+        "Go generation failed: {:?}\nstdout: {}\nstderr: {}",
+        output.status.code(),
+        String::from_utf8_lossy(&output.stdout),
+        String::from_utf8_lossy(&output.stderr),
+    );
+    assert_eq!(
+        std::fs::read_to_string(sdk.join("user-notes.txt")).unwrap(),
+        "keep me"
+    );
+    assert!(sdk.join("functions.go").is_file());
+}
+
+#[test]
+fn generate_go_removes_stale_owned_files_and_preserves_unknown_files() {
+    if !gofmt_is_available() {
+        return;
+    }
+    let built = &common::baml_cli();
+    let tmp = tempfile::tempdir().unwrap();
+    create_project_with_go_generator(
+        tmp.path(),
+        "class FormerType {\n  value: string\n}\n\nfunction echo(value: string) -> string { value }\n",
+    );
+
+    let first = run_baml_cli(built, tmp.path(), &["generate", "--from", "."]);
+    assert!(
+        first.status.success(),
+        "initial Go generation failed:\n{}",
+        String::from_utf8_lossy(&first.stderr)
+    );
+    let sdk = tmp.path().join("baml_sdk");
+    assert!(sdk.join("types.go").is_file());
+    std::fs::write(sdk.join("user-notes.txt"), "keep me").unwrap();
+
+    std::fs::write(
+        tmp.path().join("baml_src/main.baml"),
+        "function echo(value: string) -> string { value }\n",
+    )
+    .unwrap();
+    let second = run_baml_cli(built, tmp.path(), &["generate", "--from", "."]);
+    assert!(
+        second.status.success(),
+        "second Go generation failed:\n{}",
+        String::from_utf8_lossy(&second.stderr)
+    );
+    assert!(
+        !sdk.join("types.go").exists(),
+        "types.go from the removed class must not survive regeneration"
+    );
+    assert_eq!(
+        std::fs::read_to_string(sdk.join("user-notes.txt")).unwrap(),
+        "keep me"
+    );
 }
 
 // ============================================================================

baml_language/crates/baml_codegen_types/Cargo.toml

@@ -19,4 +19,13 @@ workspace = true
 [dependencies]
 baml_base = { workspace = true }
 baml_type = { workspace = true }
+serde = { workspace = true }
+serde_json = { workspace = true }
+sha2 = { workspace = true }
 strum = { workspace = true }
+
+[target.'cfg(not(target_arch = "wasm32"))'.dependencies]
+fs2 = { workspace = true }
+
+[dev-dependencies]
+tempfile = { workspace = true }