feat(functions): preserve prompt environments in pull output

Author: parkerhendo

src/functions/pull.rs

@@ -44,6 +44,8 @@ struct PullFunctionRow {
     #[serde(default)]
     function_data: Option<Value>,
     #[serde(default)]
+    environments: Option<Value>,
+    #[serde(default)]
     created: Option<String>,
     #[serde(default)]
     _xact_id: Option<String>,
@@ -64,6 +66,7 @@ struct NormalizedPrompt {
     tools: Option<Value>,
     raw_tools_json: Option<String>,
     tool_functions: Option<Value>,
+    environments: Option<Value>,
 }
 
 #[derive(Debug)]
@@ -1065,6 +1068,11 @@ fn normalize_prompt_row(row: &PullFunctionRow) -> Result<NormalizedPrompt> {
         .get("tool_functions")
         .filter(|value| !is_empty_render_value(value))
         .cloned();
+    let environments = row
+        .environments
+        .as_ref()
+        .or_else(|| prompt_data.get("environments"))
+        .and_then(normalize_environment_slugs);
     let version = row
         ._xact_id
         .as_deref()
@@ -1086,6 +1094,7 @@ fn normalize_prompt_row(row: &PullFunctionRow) -> Result<NormalizedPrompt> {
         tools,
         raw_tools_json,
         tool_functions,
+        environments,
     })
 }
 
@@ -1171,6 +1180,12 @@ fn render_project_file_ts(
                 format_ts_value(tool_functions, 2)
             ));
         }
+        if let Some(environments) = &row.environments {
+            body_lines.push(format!(
+                "  environments: {},",
+                format_ts_value(environments, 2)
+            ));
+        }
 
         out.push_str(&format!(
             "export const {var_name} = project.prompts.create({{\n"
@@ -1268,6 +1283,12 @@ fn render_project_file_py(
                 format_py_value(tool_functions, 4)
             ));
         }
+        if let Some(environments) = &row.environments {
+            out.push_str(&format!(
+                "    environments={},\n",
+                format_py_value(environments, 4)
+            ));
+        }
         out.push_str(")\n\n");
     }
 
@@ -1444,6 +1465,30 @@ fn should_unquote_object_key(key: &str) -> bool {
     chars.all(|ch| ch == '$' || ch == '_' || ch.is_ascii_alphanumeric())
 }
 
+fn normalize_environment_slugs(value: &Value) -> Option<Value> {
+    let items = value.as_array()?;
+    let mut slugs = Vec::new();
+    for item in items {
+        let slug = item
+            .as_object()
+            .and_then(|object| object.get("slug"))
+            .and_then(Value::as_str)
+            .or_else(|| item.as_str())
+            .map(str::trim)
+            .filter(|value| !value.is_empty())
+            .map(ToOwned::to_owned);
+        if let Some(slug) = slug {
+            slugs.push(Value::String(slug));
+        }
+    }
+
+    if slugs.is_empty() {
+        None
+    } else {
+        Some(Value::Array(slugs))
+    }
+}
+
 fn is_empty_render_value(value: &Value) -> bool {
     match value {
         Value::Null => true,
@@ -1600,6 +1645,7 @@ mod tests {
             description: None,
             prompt_data: None,
             function_data: None,
+            environments: None,
             created: None,
             _xact_id: None,
         };
@@ -1631,6 +1677,7 @@ mod tests {
                 description: None,
                 prompt_data: None,
                 function_data: None,
+                environments: None,
                 created: None,
                 _xact_id: None,
             },
@@ -1643,6 +1690,7 @@ mod tests {
                 description: None,
                 prompt_data: None,
                 function_data: None,
+                environments: None,
                 created: None,
                 _xact_id: None,
             },
@@ -1655,6 +1703,7 @@ mod tests {
                 description: None,
                 prompt_data: None,
                 function_data: None,
+                environments: None,
                 created: None,
                 _xact_id: None,
             },
@@ -1689,6 +1738,7 @@ mod tests {
             description: None,
             prompt_data: None,
             function_data: None,
+            environments: None,
             created: None,
             _xact_id: None,
         };
@@ -1715,6 +1765,7 @@ mod tests {
             description: None,
             prompt_data: None,
             function_data: None,
+            environments: None,
             created: None,
             _xact_id: None,
         };
@@ -1787,6 +1838,10 @@ mod tests {
                 ]
             })),
             function_data: Some(serde_json::json!({ "type": "prompt" })),
+            environments: Some(serde_json::json!([
+                { "slug": "staging" },
+                { "slug": "production" }
+            ])),
             created: None,
             _xact_id: Some("123".to_string()),
         };
@@ -1809,6 +1864,9 @@ mod tests {
         assert!(rendered.contains("    version=\"123\","));
         assert!(rendered.contains("messages=["));
         assert!(rendered.contains("model=\"gpt-4o-mini\""));
+        assert!(rendered.contains("environments=["));
+        assert!(rendered.contains("\"staging\""));
+        assert!(rendered.contains("\"production\""));
     }
 
     #[test]
@@ -1832,6 +1890,9 @@ mod tests {
                 }
             })),
             function_data: Some(serde_json::json!({ "type": "prompt" })),
+            environments: Some(serde_json::json!([
+                { "slug": "test-env" }
+            ])),
             created: None,
             _xact_id: Some("123".to_string()),
         };
@@ -1851,6 +1912,23 @@ mod tests {
         assert!(rendered.contains("export const basicMath = project.prompts.create({"));
         assert!(rendered.contains("  id: \"f1\","));
         assert!(rendered.contains("  version: \"123\","));
+        assert!(rendered.contains("environments: ["));
+        assert!(rendered.contains("\"test-env\""));
+    }
+
+    #[test]
+    fn normalize_environment_slugs_supports_objects_and_strings() {
+        let raw = serde_json::json!([
+            { "slug": " staging " },
+            "production",
+            { "slug": "" },
+            { "foo": "bar" }
+        ]);
+
+        assert_eq!(
+            normalize_environment_slugs(&raw),
+            Some(serde_json::json!(["staging", "production"]))
+        );
     }
 
     #[test]