diff --git a/src/functions/api.rs b/src/functions/api.rs index e4e1f1d..3b06132 100644 --- a/src/functions/api.rs +++ b/src/functions/api.rs @@ -1,5 +1,6 @@ use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; +use serde_json::json; use serde_json::Value; use urlencoding::encode; @@ -61,6 +62,30 @@ pub struct CodeUploadSlot { #[derive(Debug, Clone)] pub struct InsertFunctionsResult { pub ignored_entries: Option, + pub xact_id: Option, + pub functions: Vec, +} + +#[derive(Debug, Clone)] +pub struct InsertedFunction { + pub id: String, + pub slug: String, + pub project_id: String, + pub found_existing: bool, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct EnvironmentObject { + #[serde(default)] + pub environment_slug: Option, + #[serde(default)] + pub object_version: Option, +} + +#[derive(Debug, Deserialize)] +struct EnvironmentObjectsResponse { + #[serde(default)] + objects: Vec, } pub async fn list_functions( @@ -98,12 +123,7 @@ pub async fn invoke_function( client: &ApiClient, body: &serde_json::Value, ) -> Result { - let org_name = client.org_name(); - let headers = if !org_name.is_empty() { - vec![("x-bt-org-name", org_name)] - } else { - Vec::new() - }; + let headers = org_headers(client); let timeout = std::time::Duration::from_secs(300); client .post_with_headers_timeout("/function/invoke", body, &headers, Some(timeout)) @@ -223,23 +243,110 @@ pub async fn insert_functions( client: &ApiClient, functions: &[Value], ) -> Result { + let headers = org_headers(client); let body = serde_json::json!({ "functions": functions }); let raw: Value = client - .post("/insert-functions", &body) + .post_with_headers("/insert-functions", &body, &headers) .await .context("failed to insert functions")?; Ok(InsertFunctionsResult { ignored_entries: ignored_count(&raw), + xact_id: xact_id(&raw), + functions: inserted_functions(&raw), }) } +pub async fn list_environment_objects_for_prompt( + client: &ApiClient, + prompt_id: &str, +) -> Result> { + let headers = org_headers(client); + let path = format!("/environment-object/prompt/{}", encode(prompt_id)); + let response: EnvironmentObjectsResponse = client + .get_with_headers(&path, &headers) + .await + .with_context(|| format!("failed to list environments via {path}"))?; + Ok(response.objects) +} + +pub async fn upsert_environment_object_for_prompt( + client: &ApiClient, + prompt_id: &str, + environment_slug: &str, + object_version: &str, +) -> Result<()> { + let headers = org_headers(client); + let path = format!( + "/environment-object/prompt/{}/{}", + encode(prompt_id), + encode(environment_slug) + ); + let mut body = json!({ + "object_version": object_version, + }); + let org_name = client.org_name().trim(); + if !org_name.is_empty() { + body["org_name"] = Value::String(org_name.to_string()); + } + let _: Value = client + .put_with_headers(&path, &body, &headers) + .await + .with_context(|| format!("failed to upsert environment association via {path}"))?; + Ok(()) +} + +fn org_headers(client: &ApiClient) -> Vec<(&str, &str)> { + let org_name = client.org_name().trim(); + if org_name.is_empty() { + Vec::new() + } else { + vec![("x-bt-org-name", org_name)] + } +} + fn ignored_count(raw: &Value) -> Option { raw.get("ignored_count") .and_then(Value::as_u64) .and_then(|count| usize::try_from(count).ok()) } +fn xact_id(raw: &Value) -> Option { + raw.get("xact_id") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn inserted_functions(raw: &Value) -> Vec { + let Some(items) = raw.get("functions").and_then(Value::as_array) else { + return Vec::new(); + }; + + items + .iter() + .filter_map(|item| { + let id = item.get("id").and_then(Value::as_str)?.trim(); + let slug = item.get("slug").and_then(Value::as_str)?.trim(); + let project_id = item.get("project_id").and_then(Value::as_str)?.trim(); + if id.is_empty() || slug.is_empty() || project_id.is_empty() { + return None; + } + let found_existing = item + .get("found_existing") + .and_then(Value::as_bool) + .unwrap_or(false); + Some(InsertedFunction { + id: id.to_string(), + slug: slug.to_string(), + project_id: project_id.to_string(), + found_existing, + }) + }) + .collect() +} + #[cfg(test)] mod tests { use super::*; @@ -258,6 +365,29 @@ mod tests { assert_eq!(ignored_count(&serde_json::json!({})), None); } + #[test] + fn insert_functions_response_extracts_metadata() { + let raw = serde_json::json!({ + "xact_id": "123", + "functions": [ + { + "id": "fn_1", + "slug": "hello", + "project_id": "proj_1", + "found_existing": true + } + ] + }); + + assert_eq!(xact_id(&raw).as_deref(), Some("123")); + let functions = inserted_functions(&raw); + assert_eq!(functions.len(), 1); + assert_eq!(functions[0].id, "fn_1"); + assert_eq!(functions[0].slug, "hello"); + assert_eq!(functions[0].project_id, "proj_1"); + assert!(functions[0].found_existing); + } + #[test] fn parse_function_list_page_allows_non_paginated_shape() { let raw = serde_json::json!({ diff --git a/src/functions/pull.rs b/src/functions/pull.rs index 1f636cd..73ea8eb 100644 --- a/src/functions/pull.rs +++ b/src/functions/pull.rs @@ -44,6 +44,8 @@ struct PullFunctionRow { #[serde(default)] function_data: Option, #[serde(default)] + environments: Option, + #[serde(default)] created: Option, #[serde(default)] _xact_id: Option, @@ -64,6 +66,7 @@ struct NormalizedPrompt { tools: Option, raw_tools_json: Option, tool_functions: Option, + environments: Option, } #[derive(Debug)] @@ -248,6 +251,8 @@ pub async fn run(base: BaseArgs, args: PullArgs) -> Result<()> { } } + hydrate_prompt_environments(&auth_ctx.client, &mut materializable, args.verbose).await; + let output_dir = if args.output_dir.is_absolute() { args.output_dir.clone() } else { @@ -423,6 +428,78 @@ pub async fn run(base: BaseArgs, args: PullArgs) -> Result<()> { Ok(()) } +async fn hydrate_prompt_environments( + client: &crate::http::ApiClient, + rows: &mut [PullFunctionRow], + verbose: bool, +) { + for row in rows { + if row.environments.is_some() { + continue; + } + + let objects = match api::list_environment_objects_for_prompt(client, &row.id).await { + Ok(objects) => objects, + Err(err) => { + if verbose { + eprintln!( + "{} unable to load environments for '{}': {err}", + style("warning:").yellow(), + row.slug + ); + } + continue; + } + }; + + let slugs = collect_environment_slugs_for_version(&objects, row._xact_id.as_deref()); + if !slugs.is_empty() { + row.environments = Some(Value::Array(slugs.into_iter().map(Value::String).collect())); + } + } +} + +fn collect_environment_slugs_for_version( + objects: &[api::EnvironmentObject], + version: Option<&str>, +) -> Vec { + let target_version = version.map(str::trim).filter(|value| !value.is_empty()); + let target_version_num = target_version.and_then(|value| value.parse::().ok()); + let mut slugs = BTreeSet::new(); + + for object in objects { + if let Some(target_version) = target_version { + if let Some(object_version) = object + .object_version + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + { + let matches = match (target_version_num, object_version.parse::().ok()) { + (Some(expected), Some(actual)) => actual == expected, + _ => object_version == target_version, + }; + if !matches { + continue; + } + } + } + + let Some(slug) = object + .environment_slug + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + else { + continue; + }; + + slugs.insert(slug.to_string()); + } + + slugs.into_iter().collect() +} + async fn get_projects_cached<'a>( client: &crate::http::ApiClient, cache: &'a mut Option>, @@ -1065,6 +1142,11 @@ fn normalize_prompt_row(row: &PullFunctionRow) -> Result { .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 +1168,7 @@ fn normalize_prompt_row(row: &PullFunctionRow) -> Result { tools, raw_tools_json, tool_functions, + environments, }) } @@ -1171,6 +1254,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 +1357,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 +1539,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 { + 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 +1719,7 @@ mod tests { description: None, prompt_data: None, function_data: None, + environments: None, created: None, _xact_id: None, }; @@ -1631,6 +1751,7 @@ mod tests { description: None, prompt_data: None, function_data: None, + environments: None, created: None, _xact_id: None, }, @@ -1643,6 +1764,7 @@ mod tests { description: None, prompt_data: None, function_data: None, + environments: None, created: None, _xact_id: None, }, @@ -1655,6 +1777,7 @@ mod tests { description: None, prompt_data: None, function_data: None, + environments: None, created: None, _xact_id: None, }, @@ -1689,6 +1812,7 @@ mod tests { description: None, prompt_data: None, function_data: None, + environments: None, created: None, _xact_id: None, }; @@ -1715,6 +1839,7 @@ mod tests { description: None, prompt_data: None, function_data: None, + environments: None, created: None, _xact_id: None, }; @@ -1787,6 +1912,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 +1938,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 +1964,7 @@ 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 +1984,55 @@ 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] + fn collect_environment_slugs_filters_by_version() { + let objects = vec![ + api::EnvironmentObject { + environment_slug: Some("dev".to_string()), + object_version: Some("0000000000000001".to_string()), + }, + api::EnvironmentObject { + environment_slug: Some("prod".to_string()), + object_version: Some("0000000000000002".to_string()), + }, + api::EnvironmentObject { + environment_slug: Some("preview".to_string()), + object_version: None, + }, + ]; + + assert_eq!( + collect_environment_slugs_for_version(&objects, Some("0000000000000002")), + vec!["preview".to_string(), "prod".to_string()] + ); + + let padded = vec![api::EnvironmentObject { + environment_slug: Some("prod".to_string()), + object_version: Some("2".to_string()), + }]; + assert_eq!( + collect_environment_slugs_for_version(&padded, Some("0000000000000002")), + vec!["prod".to_string()] + ); } #[test] diff --git a/src/functions/push.rs b/src/functions/push.rs index 12862a3..4ad6554 100644 --- a/src/functions/push.rs +++ b/src/functions/push.rs @@ -23,6 +23,7 @@ use crate::functions::report::{ CommandStatus, FileStatus, HardFailureReason, PushFileReport, PushSummary, ReportError, SoftSkipReason, }; +use crate::http::ApiClient; use crate::js_runner; use crate::projects::api::{create_project, get_project_by_name, list_projects}; use crate::python_runner; @@ -133,6 +134,12 @@ struct FileFailure { message: String, } +#[derive(Debug, Clone)] +struct EnvironmentPlanEntry { + environment_slugs: Vec, + if_exists: String, +} + fn error_chain(err: &anyhow::Error) -> String { format!("{err:#}") } @@ -936,7 +943,7 @@ async fn push_file( }); } - let insert_result = api::insert_functions(&auth_ctx.client, &function_events) + let insert_result = insert_functions_with_environment_workaround(&auth_ctx.client, &function_events) .await .map_err(|err| FileFailure { reason: HardFailureReason::InsertFunctionsFailed, @@ -970,6 +977,194 @@ async fn push_file( }) } +async fn insert_functions_with_environment_workaround( + client: &ApiClient, + function_events: &[Value], +) -> Result { + let environment_plan = collect_environment_plan(function_events)?; + if !environment_plan.iter().any(Option::is_some) { + return api::insert_functions(client, function_events).await; + } + + let stripped_events = strip_environments_from_function_events(function_events); + let insert_result = api::insert_functions(client, &stripped_events) + .await + .context("failed to insert functions")?; + apply_environment_plan(client, &insert_result, &environment_plan).await?; + Ok(insert_result) +} + +fn collect_environment_plan( + function_events: &[Value], +) -> Result>> { + let mut plan = Vec::with_capacity(function_events.len()); + + for (index, function_event) in function_events.iter().enumerate() { + let object = function_event + .as_object() + .ok_or_else(|| anyhow!("function event entry {index} is not a JSON object"))?; + + let Some(environments) = object.get("environments") else { + plan.push(None); + continue; + }; + if environments.is_null() { + plan.push(None); + continue; + } + + let environment_list = environments.as_array().ok_or_else(|| { + anyhow!("function event entry {index} has invalid 'environments' (expected array)") + })?; + if environment_list.len() > 10 { + bail!("function event entry {index} exceeds the maximum of 10 environments"); + } + + let mut seen = BTreeSet::new(); + let mut slugs = Vec::new(); + for (env_index, env) in environment_list.iter().enumerate() { + let slug = env + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + .or_else(|| { + env.as_object() + .and_then(|value| value.get("slug")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + }) + .ok_or_else(|| { + anyhow!( + "function event entry {index} has invalid environments[{env_index}] (expected string or object with non-empty slug)" + ) + })?; + if seen.insert(slug.to_string()) { + slugs.push(slug.to_string()); + } + } + if slugs.is_empty() { + plan.push(None); + continue; + } + + let function_type = object + .get("function_data") + .and_then(Value::as_object) + .and_then(|data| data.get("type")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()); + let is_prompt_function = function_type.is_none() || function_type == Some("prompt"); + if !is_prompt_function { + let function_type = function_type.unwrap_or("unknown"); + bail!( + "environments is only supported for prompt functions, but got function_data.type=\"{function_type}\"" + ); + } + + let if_exists = object + .get("if_exists") + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .unwrap_or("error") + .to_ascii_lowercase(); + plan.push(Some(EnvironmentPlanEntry { + environment_slugs: slugs, + if_exists, + })); + } + + Ok(plan) +} + +fn strip_environments_from_function_events(function_events: &[Value]) -> Vec { + function_events + .iter() + .map(|function_event| { + let mut stripped = function_event.clone(); + if let Some(object) = stripped.as_object_mut() { + object.remove("environments"); + } + stripped + }) + .collect() +} + +async fn apply_environment_plan( + client: &ApiClient, + insert_result: &api::InsertFunctionsResult, + environment_plan: &[Option], +) -> Result<()> { + if !environment_plan.iter().any(Option::is_some) { + return Ok(()); + } + if insert_result.functions.len() != environment_plan.len() { + bail!( + "insert-functions response returned {} function records for {} inputs while applying environments", + insert_result.functions.len(), + environment_plan.len() + ); + } + + let has_updates = insert_result + .functions + .iter() + .zip(environment_plan.iter()) + .any(|(inserted_function, plan_entry)| { + plan_entry.as_ref().is_some_and(|entry| { + !entry.environment_slugs.is_empty() + && (!inserted_function.found_existing || entry.if_exists == "replace") + }) + }); + if !has_updates { + return Ok(()); + } + + let object_version = insert_result + .xact_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .ok_or_else(|| { + anyhow!("insert-functions response missing xact_id while applying environments") + })? + .to_string(); + + for (inserted_function, plan_entry) in + insert_result.functions.iter().zip(environment_plan.iter()) + { + let Some(plan_entry) = plan_entry else { + continue; + }; + if inserted_function.found_existing && plan_entry.if_exists != "replace" { + continue; + } + + for environment_slug in &plan_entry.environment_slugs { + api::upsert_environment_object_for_prompt( + client, + &inserted_function.id, + environment_slug, + &object_version, + ) + .await + .with_context(|| { + format!( + "failed to upsert environment association for prompt '{}' ({}) in project '{}' and environment '{}'", + inserted_function.slug, + inserted_function.id, + inserted_function.project_id, + environment_slug + ) + })?; + } + } + + Ok(()) +} + fn build_js_bundle( source_path: &Path, args: &PushArgs, @@ -3030,6 +3225,61 @@ mod tests { assert_eq!(calculate_upload_counts(3, None), (3, 0)); } + #[test] + fn environment_plan_collects_prompt_environments() { + let function_events = vec![serde_json::json!({ + "project_id": "proj_1", + "slug": "hello", + "if_exists": "replace", + "function_data": { "type": "prompt" }, + "environments": [{ "slug": "staging" }, "prod"], + })]; + + let plan = collect_environment_plan(&function_events).expect("collect plan"); + assert_eq!(plan.len(), 1); + let entry = plan[0].as_ref().expect("plan entry"); + assert_eq!(entry.environment_slugs, vec!["staging", "prod"]); + assert_eq!(entry.if_exists, "replace"); + } + + #[test] + fn environment_plan_rejects_non_prompt_function_type() { + let function_events = vec![serde_json::json!({ + "project_id": "proj_1", + "slug": "hello", + "if_exists": "replace", + "function_data": { "type": "code" }, + "environments": [{ "slug": "staging" }], + })]; + + let err = collect_environment_plan(&function_events).expect_err("must fail"); + assert!(err + .to_string() + .contains("environments is only supported for prompt functions")); + } + + #[test] + fn strip_environments_removes_environment_field() { + let function_events = vec![serde_json::json!({ + "project_id": "proj_1", + "slug": "hello", + "environments": [{ "slug": "staging" }], + "metadata": { "source": "test" } + })]; + + let stripped = strip_environments_from_function_events(&function_events); + assert_eq!(stripped.len(), 1); + let obj = stripped[0].as_object().expect("object"); + assert!(!obj.contains_key("environments")); + assert_eq!( + obj.get("metadata") + .and_then(Value::as_object) + .and_then(|metadata| metadata.get("source")) + .and_then(Value::as_str), + Some("test") + ); + } + #[test] fn requirements_reference_escape_is_rejected() { let dir = tempfile::tempdir().expect("tempdir"); diff --git a/src/http.rs b/src/http.rs index e077add..9d26497 100644 --- a/src/http.rs +++ b/src/http.rs @@ -168,6 +168,57 @@ impl ApiClient { request.send().await.context("request failed") } + pub async fn get_with_headers( + &self, + path: &str, + headers: &[(&str, &str)], + ) -> Result { + let url = self.url(path); + let mut request = self.http.get(&url).bearer_auth(&self.api_key); + + for (key, value) in headers { + request = request.header(*key, *value); + } + + let response = request.send().await.context("request failed")?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(HttpError { status, body }.into()); + } + + response.json().await.context("failed to parse response") + } + + pub async fn put_with_headers( + &self, + path: &str, + body: &B, + headers: &[(&str, &str)], + ) -> Result + where + T: DeserializeOwned, + B: Serialize, + { + let url = self.url(path); + let mut request = self.http.put(&url).bearer_auth(&self.api_key).json(body); + + for (key, value) in headers { + request = request.header(*key, *value); + } + + let response = request.send().await.context("request failed")?; + + if !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(HttpError { status, body }.into()); + } + + response.json().await.context("failed to parse response") + } + pub async fn delete(&self, path: &str) -> Result<()> { let url = self.url(path); let response = self diff --git a/tests/functions.rs b/tests/functions.rs index 9347b4f..dfcaab1 100644 --- a/tests/functions.rs +++ b/tests/functions.rs @@ -200,6 +200,8 @@ struct MockServerState { requests: Mutex>, projects: Mutex>, pull_rows: Mutex>, + environment_objects: Mutex>>, + environment_upserts: Mutex>, uploaded_bundles: Mutex>>, inserted_functions: Mutex>, bundle_counter: Mutex, @@ -227,6 +229,14 @@ impl MockServer { .route("/upload/{bundle_id}", web::put().to(mock_upload_bundle)) .route("/insert-functions", web::post().to(mock_insert_functions)) .route("/v1/function", web::get().to(mock_list_functions)) + .route( + "/environment-object/{object_type}/{object_id}", + web::get().to(mock_list_environment_objects), + ) + .route( + "/environment-object/{object_type}/{object_id}/{environment_slug}", + web::put().to(mock_upsert_environment_object), + ) }) .workers(1) .listen(listener) @@ -363,7 +373,33 @@ async fn mock_insert_functions( .lock() .expect("inserted functions lock"); inserted.extend(body.functions.clone()); - HttpResponse::Ok().json(serde_json::json!({ "ignored_count": 0 })) + let functions = body + .functions + .iter() + .enumerate() + .map(|(index, function)| { + let slug = function + .get("slug") + .and_then(Value::as_str) + .unwrap_or("unknown"); + let project_id = function + .get("project_id") + .and_then(Value::as_str) + .unwrap_or("proj_mock"); + serde_json::json!({ + "id": format!("fn_inserted_{index}"), + "slug": slug, + "project_id": project_id, + "found_existing": false + }) + }) + .collect::>(); + + HttpResponse::Ok().json(serde_json::json!({ + "ignored_count": 0, + "xact_id": "0000000000000001", + "functions": functions + })) } async fn mock_list_functions( @@ -399,6 +435,77 @@ async fn mock_list_functions( })) } +async fn mock_list_environment_objects( + state: web::Data>, + req: HttpRequest, +) -> HttpResponse { + log_request(&state, &req); + let object_type = req.match_info().get("object_type").unwrap_or_default(); + if object_type != "prompt" { + return HttpResponse::Ok().json(serde_json::json!({ "objects": [] })); + } + let object_id = req.match_info().get("object_id").unwrap_or_default(); + let rows = state + .environment_objects + .lock() + .expect("environment objects lock") + .get(object_id) + .cloned() + .unwrap_or_default(); + + HttpResponse::Ok().json(serde_json::json!({ "objects": rows })) +} + +async fn mock_upsert_environment_object( + state: web::Data>, + req: HttpRequest, + body: web::Json, +) -> HttpResponse { + log_request(&state, &req); + let object_type = req.match_info().get("object_type").unwrap_or_default(); + let object_id = req.match_info().get("object_id").unwrap_or_default(); + let environment_slug = req.match_info().get("environment_slug").unwrap_or_default(); + let object_version = body + .get("object_version") + .and_then(Value::as_str) + .unwrap_or_default(); + let org_name = body + .get("org_name") + .and_then(Value::as_str) + .unwrap_or_default(); + + state + .environment_upserts + .lock() + .expect("environment upserts lock") + .push(serde_json::json!({ + "object_type": object_type, + "object_id": object_id, + "environment_slug": environment_slug, + "object_version": object_version, + "org_name": org_name + })); + + if object_type == "prompt" + && !object_id.trim().is_empty() + && !environment_slug.trim().is_empty() + && !object_version.trim().is_empty() + { + state + .environment_objects + .lock() + .expect("environment objects lock") + .entry(object_id.to_string()) + .or_default() + .push(serde_json::json!({ + "environment_slug": environment_slug, + "object_version": object_version + })); + } + + HttpResponse::Ok().json(serde_json::json!({})) +} + fn log_request(state: &Arc, req: &HttpRequest) { let entry = if req.query_string().is_empty() { req.path().to_string() @@ -1722,6 +1829,199 @@ exit 24 ); } +#[cfg(unix)] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn functions_push_prompt_environments_upsert_after_insert() { + if !command_exists("node") { + eprintln!( + "Skipping functions_push_prompt_environments_upsert_after_insert (node not installed)." + ); + return; + } + + let state = Arc::new(MockServerState::default()); + state + .projects + .lock() + .expect("projects lock") + .push(MockProject { + id: "proj_mock".to_string(), + name: "mock-project".to_string(), + org_id: "org_mock".to_string(), + }); + let server = MockServer::start(state.clone()).await; + + let tmp = tempdir().expect("tempdir"); + let source = tmp.path().join("prompt.js"); + std::fs::write( + &source, + "globalThis._evals ??= { functions: [], prompts: [], parameters: [], evaluators: {}, reporters: {} };\n", + ) + .expect("write source file"); + + let runner = tmp.path().join("mock-runner.sh"); + std::fs::write( + &runner, + r#"#!/bin/sh +set -eu +_runner_script="$1" +shift +_runner_name="$(basename "$_runner_script")" + +if [ "$_runner_name" = "functions-runner.ts" ]; then +node - "$@" <<'NODE' +const path = require("node:path"); +const files = process.argv.slice(2); +const manifest = { + runtime_context: { runtime: "node", version: process.versions.node || "unknown" }, + files: files.map((file) => ({ + source_file: path.resolve(file), + entries: [ + { + kind: "function_event", + project_id: "proj_mock", + event: { + project_id: "proj_mock", + name: "mock-prompt", + slug: "mock-prompt", + function_data: { type: "prompt" }, + prompt_data: { + prompt: { + type: "chat", + messages: [{ role: "system", content: "Hello" }] + }, + options: { model: "gpt-4o-mini" } + }, + if_exists: "replace", + environments: [{ slug: "staging" }, "prod"] + } + } + ] + })) +}; +process.stdout.write(JSON.stringify(manifest)); +NODE +exit 0 +fi + +if [ "$_runner_name" = "functions-bundler.ts" ]; then + _source_file="$1" + _output_file="$2" + cp "$_source_file" "$_output_file" + exit 0 +fi + +echo "unexpected runner script: $_runner_name" >&2 +exit 24 +"#, + ) + .expect("write mock runner"); + use std::os::unix::fs::PermissionsExt; + let mut perms = std::fs::metadata(&runner) + .expect("runner metadata") + .permissions(); + perms.set_mode(0o755); + std::fs::set_permissions(&runner, perms).expect("runner permissions"); + + let output = Command::new(bt_binary_path()) + .current_dir(tmp.path()) + .args([ + "functions", + "--json", + "push", + "--file", + source + .to_str() + .expect("source path should be valid UTF-8 for test"), + "--language", + "javascript", + "--runner", + runner + .to_str() + .expect("runner path should be valid UTF-8 for test"), + "--if-exists", + "replace", + ]) + .env("BRAINTRUST_API_KEY", "test-key") + .env("BRAINTRUST_ORG_NAME", "test-org") + .env("BRAINTRUST_API_URL", &server.base_url) + .env("BRAINTRUST_APP_URL", &server.base_url) + .env("BRAINTRUST_NO_COLOR", "1") + .env_remove("BRAINTRUST_PROFILE") + .output() + .expect("run bt functions push"); + + server.stop().await; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + panic!("mock push failed:\n{stderr}"); + } + + let summary: Value = serde_json::from_slice(&output.stdout).expect("parse push summary"); + assert_eq!(summary["status"].as_str(), Some("success")); + assert_eq!(summary["uploaded_files"].as_u64(), Some(1)); + assert_eq!(summary["failed_files"].as_u64(), Some(0)); + + let inserted = state + .inserted_functions + .lock() + .expect("inserted functions lock") + .clone(); + assert_eq!(inserted.len(), 1, "exactly one function should be inserted"); + let first = inserted[0].as_object().expect("inserted function object"); + assert!( + first.get("environments").is_none(), + "insert payload should strip environments and use environment-object endpoint" + ); + + let upserts = state + .environment_upserts + .lock() + .expect("environment upserts lock") + .clone(); + assert_eq!(upserts.len(), 2, "expected one upsert per environment"); + let mut slugs = upserts + .iter() + .filter_map(|entry| entry.get("environment_slug").and_then(Value::as_str)) + .map(ToOwned::to_owned) + .collect::>(); + slugs.sort(); + assert_eq!(slugs, vec!["prod".to_string(), "staging".to_string()]); + for upsert in &upserts { + assert_eq!( + upsert.get("object_type").and_then(Value::as_str), + Some("prompt") + ); + assert_eq!( + upsert.get("object_id").and_then(Value::as_str), + Some("fn_inserted_0") + ); + assert_eq!( + upsert.get("object_version").and_then(Value::as_str), + Some("0000000000000001") + ); + assert_eq!( + upsert.get("org_name").and_then(Value::as_str), + Some("test-org") + ); + } + + let requests = state.requests.lock().expect("requests lock").clone(); + assert!( + requests + .iter() + .any(|entry| entry == "/environment-object/prompt/fn_inserted_0/staging"), + "push should upsert staging environment for the inserted prompt" + ); + assert!( + requests + .iter() + .any(|entry| entry == "/environment-object/prompt/fn_inserted_0/prod"), + "push should upsert prod environment for the inserted prompt" + ); +} + #[cfg(unix)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn functions_push_external_packages_bundles_with_runner() { @@ -2182,6 +2482,23 @@ async fn functions_pull_works_against_mock_api() { }, "_xact_id": "0000000000000001" })); + state + .environment_objects + .lock() + .expect("environment objects lock") + .insert( + "fn_123".to_string(), + vec![ + serde_json::json!({ + "environment_slug": "staging", + "object_version": "0000000000000001" + }), + serde_json::json!({ + "environment_slug": "prod", + "object_version": "0000000000000000" + }), + ], + ); let server = MockServer::start(state.clone()).await; @@ -2243,6 +2560,18 @@ async fn functions_pull_works_against_mock_api() { rendered.contains("gpt-4o-mini"), "rendered file should include model config" ); + assert!( + rendered.contains("environments: ["), + "rendered file should include environments field" + ); + assert!( + rendered.contains("\"staging\""), + "rendered file should include matching environment slug" + ); + assert!( + !rendered.contains("\"prod\""), + "rendered file should omit non-matching environment versions" + ); let requests = state.requests.lock().expect("requests lock").clone(); assert!( @@ -2253,6 +2582,12 @@ async fn functions_pull_works_against_mock_api() { }), "pull request should include selector query params" ); + assert!( + requests + .iter() + .any(|entry| entry == "/environment-object/prompt/fn_123"), + "pull should hydrate prompt environments from environment-object endpoint" + ); } #[tokio::test(flavor = "multi_thread", worker_threads = 2)]