craft-ai-core: track model provenance by address; require the inputs a model needs

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<ModelInfo>,
     /// Editing: the saved key's locator, and whether a key was saved at all.
     secret: Option<SecretRef>,
-    original_endpoint: Option<String>,
+    /// The address `models` were loaded, preset or last set for.
-    /// Editing: the models the connection had when the form opened, which belong to the original
-    /// address.
-    loaded_models: Vec<ModelInfo>,
-    /// The address `models` were last set for through [`set_models`](Self::set_models).
     models_for: Option<String>,
 }
 
@@ -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(),
+            models_for: c.endpoint.clone(),
-            loaded_models: c.models.clone(),
-            models_for: None,
         }
     }
 
@@ -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
+    /// Known model provenance must match the current address, including after repeated edits.
-    /// address's: neither found by a test at this address nor different from what was loaded.
+    /// 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);
+        self.models_for.as_deref().is_none_or(|origin| origin == address)
-        !moved || self.models_for.as_deref() == Some(address) || (self.models_for.is_none() && self.models != self.loaded_models)
     }
 
     /// Where this would send data, as far as the address can tell. `None` while it is not a usable

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);
 }
 

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<InputRole>,
 }
@@ -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)];

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<ModelInfo> {
+    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());
+}