diff --git a/crates/ai-core/src/draft.rs b/crates/ai-core/src/draft.rs index 3cd9cbf6..bd70a4a0 100644 --- a/crates/ai-core/src/draft.rs +++ b/crates/ai-core/src/draft.rs @@ -244,15 +244,12 @@ pub struct ConnectionDraft { /// Header that carries the key; empty sends it as `Authorization: Bearer`. pub key_header: String, pub enabled: bool, - /// Models found by a test, carried onto the saved connection. + /// Models carried onto the saved connection. Use [`set_models`](Self::set_models) when + /// replacing models after an address change so their address is recorded. pub models: Vec, /// Editing: the saved key's locator, and whether a key was saved at all. secret: Option, - original_endpoint: Option, - /// Editing: the models the connection had when the form opened, which belong to the original - /// address. - loaded_models: Vec, - /// The address `models` were last set for through [`set_models`](Self::set_models). + /// The address `models` were loaded, preset or last set for. models_for: Option, } @@ -270,8 +267,6 @@ impl Default for ConnectionDraft { enabled: true, models: Vec::new(), secret: None, - original_endpoint: None, - loaded_models: Vec::new(), models_for: None, } } @@ -294,13 +289,15 @@ impl fmt::Debug for ConnectionDraft { impl ConnectionDraft { /// A new connection started from a preset. pub fn from_preset(preset: &Preset) -> Self { + let models = preset_models(preset); Self { label: preset.name.to_owned(), provider: preset.provider.to_owned(), endpoint: preset.endpoint.to_owned(), billing: preset.billing, uses_key: preset.needs_key, - models: preset_models(preset), + models_for: (!models.is_empty()).then(|| preset.endpoint.trim().to_owned()), + models, ..Self::default() } } @@ -319,9 +316,7 @@ impl ConnectionDraft { enabled: c.enabled, models: c.models.clone(), secret: c.auth.secret_ref().cloned(), - original_endpoint: c.endpoint.clone(), - loaded_models: c.models.clone(), - models_for: None, + models_for: c.endpoint.clone(), } } @@ -332,11 +327,10 @@ impl ConnectionDraft { self.models_for = Some(self.endpoint.trim().to_owned()); } - /// False when the address moved since the form opened and `models` are still the old - /// address's: neither found by a test at this address nor different from what was loaded. + /// Known model provenance must match the current address, including after repeated edits. + /// A new custom draft may also have manually declared models without a prior address. fn models_fit(&self, address: &str) -> bool { - let moved = self.original_endpoint.as_deref().is_some_and(|o| o != address); - !moved || self.models_for.as_deref() == Some(address) || (self.models_for.is_none() && self.models != self.loaded_models) + self.models_for.as_deref().is_none_or(|origin| origin == address) } /// Where this would send data, as far as the address can tell. `None` while it is not a usable diff --git a/crates/ai-core/src/draft/tests.rs b/crates/ai-core/src/draft/tests.rs index 20e26280..7ac10cc5 100644 --- a/crates/ai-core/src/draft/tests.rs +++ b/crates/ai-core/src/draft/tests.rs @@ -180,7 +180,7 @@ fn models_found_at_the_new_address_are_kept() { let found = vec![ModelInfo::new(crate::ids::ModelId::new("llama-x").unwrap(), "Llama X")]; let mut moved = ConnectionDraft::from_connection(&saved); moved.endpoint = "https://other.example.com/v1".into(); - moved.models = found.clone(); + moved.set_models(found.clone()); assert_eq!(moved.prepare(std::slice::from_ref(&saved)).unwrap().connection.models, found); } diff --git a/crates/ai-core/src/role.rs b/crates/ai-core/src/role.rs index 2b1d7f06..4a70cef7 100644 --- a/crates/ai-core/src/role.rs +++ b/crates/ai-core/src/role.rs @@ -22,7 +22,7 @@ pub struct Role { pub id: RoleId, pub label: String, pub capability: CapabilityId, - /// Inputs this feature will send, so a model that cannot take them is not offered. + /// Inputs this feature will send. A model must accept them and require no others. #[serde(default)] pub inputs: Vec, } @@ -52,6 +52,7 @@ pub enum ResolveError { ModelMissing { connection: ConnectionId, model: ModelId }, CapabilityMismatch { model: ModelId, capability: CapabilityId }, InputUnsupported { model: ModelId, input: InputRole }, + InputRequired { model: ModelId, input: InputRole }, ToolMismatch(RoleId), } @@ -70,6 +71,9 @@ impl std::fmt::Display for ResolveError { Self::InputUnsupported { model, input } => { write!(f, "model {model} does not accept a '{input}' input") } + Self::InputRequired { model, input } => { + write!(f, "model {model} requires a '{input}' input that this tool does not provide") + } Self::ToolMismatch(id) => write!(f, "tool {id} does not match its shared definition"), } } @@ -94,14 +98,18 @@ pub fn resolve<'a>(role: &Role, layers: &[&[Binding]], connections: &'a [Connect if let Some(input) = role.inputs.iter().find(|i| !spec.accepts.contains_key(i)) { return Err(ResolveError::InputUnsupported { model: model.id.clone(), input: input.clone() }); } + if let Some(input) = spec.requires.iter().find(|i| !role.inputs.contains(i)) { + return Err(ResolveError::InputRequired { model: model.id.clone(), input: input.clone() }); + } Ok(Resolved { connection, model, spec }) } /// What `model` can do for `role`: its spec for the role's capability, if it takes every input the -/// role sends. `None` means the model cannot fill the role. +/// role sends and the role supplies every input the model requires. `None` means the model +/// cannot fill the role. pub fn fits<'a>(role: &Role, model: &'a ModelInfo) -> Option<&'a CapabilitySpec> { let spec = model.spec_for(&role.capability)?; - role.inputs.iter().all(|i| spec.accepts.contains_key(i)).then_some(spec) + (role.inputs.iter().all(|i| spec.accepts.contains_key(i)) && spec.requires.iter().all(|i| role.inputs.contains(i))).then_some(spec) } /// Every enabled connection's models that could fill `role`. This is the model picker's list. @@ -187,6 +195,17 @@ mod tests { assert!(matches!(err, ResolveError::InputUnsupported { input: InputRole::Mask, .. })); } + #[test] + fn a_required_input_the_role_cannot_supply_has_a_specific_error() { + let mut conns = [conn("a", true)]; + conns[0].models[0].capabilities[0].requires = vec![InputRole::Mask]; + let mut r = role(); + r.inputs = vec![InputRole::Source]; + let err = resolve(&r, &[&[bind("a")]], &conns).unwrap_err(); + assert!(matches!(err, ResolveError::InputRequired { input: InputRole::Mask, .. })); + assert_eq!(err.to_string(), "model m requires a 'mask' input that this tool does not provide"); + } + #[test] fn wrong_capability_and_missing_model() { let mut conns = [conn("a", true)]; diff --git a/crates/ai-core/tests/contract_regressions.rs b/crates/ai-core/tests/contract_regressions.rs new file mode 100644 index 00000000..871ae95d --- /dev/null +++ b/crates/ai-core/tests/contract_regressions.rs @@ -0,0 +1,135 @@ +use craft_ai_core::role::{eligible, fits, resolve}; +use craft_ai_core::*; + +fn draft(endpoint: &str) -> ConnectionDraft { + let mut draft = ConnectionDraft::default(); + draft.label = "Test".into(); + draft.endpoint = endpoint.into(); + draft.uses_key = false; + draft +} + +fn models(id: &str) -> Vec { + vec![ModelInfo::new(ModelId::new(id).unwrap(), id)] +} + +#[test] +fn new_draft_discards_models_when_address_changes_after_probe() { + let mut draft = draft("https://a.example/v1"); + draft.set_models(models("a-only")); + draft.endpoint = "https://b.example/v1".into(); + assert!(draft.prepare(&[]).unwrap().connection.models.is_empty()); +} + +#[test] +fn returning_to_original_address_does_not_keep_another_addresses_models() { + let saved = draft("https://a.example/v1").prepare(&[]).unwrap().connection; + let mut draft = ConnectionDraft::from_connection(&saved); + draft.endpoint = "https://b.example/v1".into(); + draft.set_models(models("b-only")); + draft.endpoint = "https://a.example/v1".into(); + assert!(draft.prepare(&[saved]).unwrap().connection.models.is_empty()); +} + +#[test] +fn changing_model_contents_does_not_establish_address_provenance() { + let saved = draft("https://a.example/v1").prepare(&[]).unwrap().connection; + let mut draft = ConnectionDraft::from_connection(&saved); + draft.models = models("refreshed-at-a"); + draft.endpoint = "https://b.example/v1".into(); + assert!(draft.prepare(&[saved]).unwrap().connection.models.is_empty()); +} + +#[test] +fn preset_models_belong_to_the_preset_address() { + let preset = PRESETS.iter().find(|preset| preset.key == "vertex-express").unwrap(); + let mut draft = ConnectionDraft::from_preset(preset); + draft.key = "test-key".into(); + assert!(!draft.prepare(&[]).unwrap().connection.models.is_empty()); + draft.endpoint = "https://other.example/v1".into(); + assert!(draft.prepare(&[]).unwrap().connection.models.is_empty()); +} + +#[test] +fn a_fresh_test_at_each_address_keeps_its_models() { + let saved = draft("https://a.example/v1").prepare(&[]).unwrap().connection; + let mut draft = ConnectionDraft::from_connection(&saved); + for address in ["https://b.example/v1", "https://a.example/v1", "https://c.example/v1"] { + draft.endpoint = format!(" {address} "); + draft.set_models(models("same-model")); + assert_eq!(draft.prepare(std::slice::from_ref(&saved)).unwrap().connection.models, models("same-model")); + } +} + +#[test] +fn editing_and_returning_without_a_probe_preserves_original_models() { + let mut saved = draft("https://a.example/v1").prepare(&[]).unwrap().connection; + saved.models = models("a-only"); + let mut draft = ConnectionDraft::from_connection(&saved); + draft.endpoint = "https://b.example/v1".into(); + assert!(draft.prepare(std::slice::from_ref(&saved)).unwrap().connection.models.is_empty()); + draft.endpoint = "https://a.example/v1".into(); + assert_eq!(draft.prepare(std::slice::from_ref(&saved)).unwrap().connection.models, saved.models); +} + +#[test] +fn custom_manual_models_still_work_for_a_new_connection_and_unchanged_address() { + let mut draft = ConnectionDraft::custom(); + draft.endpoint = "https://custom.example/v1".into(); + draft.uses_key = false; + draft.models = models("manual-model"); + let saved = draft.prepare(&[]).unwrap().connection; + assert_eq!(saved.models, models("manual-model")); + let mut edited = ConnectionDraft::from_connection(&saved); + edited.models = models("other-manual-model"); + assert_eq!(edited.prepare(&[saved]).unwrap().connection.models, models("other-manual-model")); +} + +fn role_and_connection() -> (Role, Connection, Binding) { + let capability = CapabilityId::parse("image.inpaint").unwrap(); + let role = Role { + id: RoleId::new("source-only").unwrap(), + label: "Source only".into(), + capability: capability.clone(), + inputs: vec![InputRole::Source], + }; + let mut connection = draft("https://a.example/v1").prepare(&[]).unwrap().connection; + connection.models = vec![ + ModelInfo::new(ModelId::new("needs-mask").unwrap(), "Needs mask").with( + CapabilitySpec::new(capability) + .accepting(InputRole::Source, 1) + .accepting(InputRole::Mask, 1) + .requiring(InputRole::Source) + .requiring(InputRole::Mask), + ), + ]; + let binding = Binding { role: role.id.clone(), connection: connection.id.clone(), model: connection.models[0].id.clone() }; + (role, connection, binding) +} + +#[test] +fn model_picker_excludes_inputs_the_tool_cannot_supply() { + let (role, connection, _) = role_and_connection(); + let request = Request::new(role.capability.clone()).with_input(InputRole::Source, Payload::inline("image/png", vec![1]).unwrap()); + assert!(request.validate(&connection.models[0].capabilities[0], &Limits::default()).is_err()); + assert!(fits(&role, &connection.models[0]).is_none()); + assert!(eligible(&role, &[connection]).is_empty()); +} + +#[test] +fn binding_resolution_rejects_inputs_the_tool_cannot_supply() { + let (role, connection, binding) = role_and_connection(); + assert!(resolve(&role, &[&[binding]], &[connection]).is_err()); +} + +#[test] +fn required_inputs_resolve_when_supplied_and_optional_inputs_need_not_be_sent() { + let (mut role, mut connection, binding) = role_and_connection(); + role.inputs.push(InputRole::Mask); + assert!(fits(&role, &connection.models[0]).is_some()); + assert!(resolve(&role, &[std::slice::from_ref(&binding)], std::slice::from_ref(&connection)).is_ok()); + role.inputs = vec![InputRole::Source]; + connection.models[0].capabilities[0].requires = vec![InputRole::Source]; + assert!(fits(&role, &connection.models[0]).is_some()); + assert!(resolve(&role, &[&[binding]], &[connection]).is_ok()); +}