Handle attributes and self-closing tags

lib/xml_schema_validator/Cargo.toml

@@ -10,7 +10,8 @@ name = "xml_schema_validator"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
-pyo3 = { version = "0.24", features = ["extension-module"] }
+pyo3 = { version = "0.24", features = ["extension-module", "auto-initialize"] }
 quick-xml = { version = "0.28.1", features = ["serde", "serialize"] }
 serde = { version = "1", features = ["derive"] }
+serde_json = "1.0"
 thiserror = "1.0.38"

lib/xml_schema_validator/README.md

@@ -1,3 +1,3 @@
 # XML Schema Validator
 
-XML Schema validation library optimized for LLM token validation.
+An XML Schema validation library optimized for LLM token validation. This library allows you to validate XML tokens against an XSD schema, enabling constrained token generation in language models.

lib/xml_schema_validator/python/xml_schema_validator/__init__.py

@@ -12,22 +12,31 @@ class XmlLogitsProcessor(LogitsProcessor):
     by setting their logits to negative infinity.
     """
     
-    def __init__(self, tokenizer: AutoTokenizer, schema_text: str):
+    def __init__(self, tokenizer: AutoTokenizer, schema_text: str = None, core=None):
         """
         Initialize the processor with a schema and tokenizer.
         
         Args:
             tokenizer: The tokenizer to use for decoding tokens
             schema_text: The XSD schema text to validate against
+            core: An existing core processor (for internal use in copy())
         """
         self.eos_token_id = tokenizer.eos_token_id
         self.tokenizer = tokenizer
-        vocab = tokenizer.get_vocab()  # This is {token: id}
-        items = dict()
-        for token, id in vocab.items():
-            items[id] = token
-        if schema_text:
+        
+        if core is not None:
+            # Used for copy() operation
+            self.core = core
+        elif schema_text:
+            # Normal initialization
+            vocab = tokenizer.get_vocab()  # This is {token: id}
+            items = dict()
+            for token, id in vocab.items():
+                items[id] = token
             self.core = XmlLogitsProcessorCore(items, schema_text)
+        else:
+            raise ValueError("Either schema_text or core must be provided")
+            
         self.prompt_length = None
         self.is_first_call = True
 
@@ -60,7 +69,10 @@ class XmlLogitsProcessor(LogitsProcessor):
             batch_processor = self.core.copy()
 
             if generated_text:
-                batch_processor.append(generated_text) 
+                # Try to append the current text to the validator
+                append_result = batch_processor.append(generated_text) 
+                
+                # If we've reached a valid EOF state, only allow the EOS token
                 if batch_processor.eof():
                     vocab_size = scores.shape[-1]
                     valid_tokens_mask = torch.zeros(vocab_size, dtype=torch.bool, device=scores.device)
@@ -68,8 +80,17 @@ class XmlLogitsProcessor(LogitsProcessor):
                     invalid_tokens_mask = ~valid_tokens_mask
                     scores[batch_idx, invalid_tokens_mask] = float('-inf')
                     continue
+                
+                # If the validation failed altogether, this is an invalid path
+                if not append_result:
+                    # Allow only EOS token if validation fails
+                    scores[batch_idx, :] = float('-inf')
+                    scores[batch_idx, self.eos_token_id] = 0
+                    continue
 
-            for token in batch_processor.get_invalid_tokens():
+            # Get tokens that would lead to invalid XML and mask them
+            invalid_tokens = batch_processor.get_invalid_tokens()
+            for token in invalid_tokens:
                 if token < scores.shape[1]:
                     scores[batch_idx, token] = float('-inf')
 
@@ -82,9 +103,9 @@ class XmlLogitsProcessor(LogitsProcessor):
         Returns:
             A new instance of the processor
         """
-        cloned = XmlLogitsProcessor(self.tokenizer, None)
-        cloned.eos_token_id = self.eos_token_id
-        cloned.core = self.core.copy()
+        # Create a new instance using the existing core
+        cloned_core = self.core.copy()
+        cloned = XmlLogitsProcessor(self.tokenizer, core=cloned_core)
         cloned.prompt_length = self.prompt_length
         cloned.is_first_call = self.is_first_call
         return cloned

lib/xml_schema_validator/src/lib.rs

@@ -18,6 +18,7 @@ use token::*;
 #[derive(Clone, Debug)]
 pub struct XmlSchemaValidator {
     current_tokens: Vec<Token>,
+    schema: Arc<schema::XsSchema>,
 }
 
 /// The Result type used throughout this crate
@@ -38,6 +39,7 @@ impl XmlSchemaValidator {
         
         Ok(Self {
             current_tokens,
+            schema: schema_arc,
         })
     }
     
@@ -51,6 +53,7 @@ impl XmlSchemaValidator {
 
     /// Check if the validator has reached the end of the XML
     pub fn eof(&self) -> bool {
+        println!("EOF: {:?}", self.current_tokens);
         self.current_tokens.is_empty() || self.current_tokens.iter().all(Token::is_eof)
     }
 
@@ -61,11 +64,25 @@ impl XmlSchemaValidator {
             new_tokens.append(&mut token.append(c));
         }
         if new_tokens.is_empty() {
-            return Err(Error::InvalidXml("No continuations".to_string()));
+            return Err(Error::InvalidXml(format!("No valid continuations for character: '{}'", c)));
         }
         self.current_tokens = new_tokens;
         Ok(())
     }
+    
+    /// Try to validate a continuation character
+    pub fn can_continue_with(&self, c: char) -> bool {
+        // Clone the current state and try to append the character
+        let mut validator_clone = self.clone();
+        validator_clone.append_char(c).is_ok()
+    }
+    
+    /// Get a list of all element names in the schema
+    pub fn get_element_names(&self) -> Vec<String> {
+        self.schema.xs_element.iter()
+            .map(|element| element.name.clone())
+            .collect()
+    }
 }
 
 #[cfg(test)]
@@ -111,4 +128,472 @@ mod tests {
         let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
         validator.append(input).expect_err("Should fail with mismatched closing tag");
     }
+    
+    #[test]
+    fn test_delete_with_attribute() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        let input = r#"<delete id="1234567890"/>"#;
+        let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        validator.append(input).unwrap();
+        assert!(validator.eof());
+    }
+    
+    #[test]
+    fn test_delete_missing_required_attribute() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        // Try to validate an incomplete delete tag that's missing the required id attribute
+        let input = "<delete></delete>";
+        let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        let result = validator.append(input);
+        assert!(result.is_err(), "Should fail without required id attribute");
+        // Check the error message explicitly
+        if let Err(err) = result {
+            println!("Error message: {}", err);
+            assert!(err.to_string().contains("required"), 
+                    "Error should mention missing required attribute");
+        }
+    }
+    
+    #[test]
+    fn test_stop_self_closing() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        let input = "<stop/>";
+        let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        validator.append(input).unwrap();
+        assert!(validator.eof());
+    }
+    
+    #[test]
+    fn test_read_stdin_self_closing() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        let input = "<read_stdin/>";
+        let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        validator.append(input).unwrap();
+        assert!(validator.eof());
+    }
+    
+    #[test]
+    fn test_single_with_attributes() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        let input = r#"<single timeout="2.5" limit="2048">ls -la</single>"#;
+        let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        validator.append(input).unwrap();
+        assert!(validator.eof());
+    }
+    
+    #[test]
+    fn test_reasoning_with_content() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        let input = "<reasoning>I should explore the file system for interesting files.</reasoning>";
+        let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        validator.append(input).unwrap();
+        assert!(validator.eof());
+    }
+    
+    #[test]
+    fn test_write_stdout_with_content() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        let input = "<write_stdout>Hello world!</write_stdout>";
+        let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        validator.append(input).unwrap();
+        assert!(validator.eof());
+    }
+
+    #[test]
+    fn test_create_validator() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        let validator = XmlSchemaValidator::new(&schema_text);
+        assert!(validator.is_ok(), "Failed to create validator: {:?}", validator.err());
+    }
+
+    #[test]
+    fn test_get_element_names() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        let validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        let names = validator.get_element_names();
+        assert!(names.contains(&"reasoning".to_string()), 
+                "Element names should contain 'reasoning', got: {:?}", names);
+    }
+
+    // Test basic XML fragment validation - debugging the "<reasoning>" failure
+    #[test]
+    fn test_append_reasoning_tag() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        let validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        // Test character by character for precise failure point
+        assert!(validator.can_continue_with('<'), "Should accept '<'");
+        // Clone at each step to debug exact point of failure
+        let mut v1 = validator.clone();
+        assert!(v1.append_char('<').is_ok(), "Failed to append '<'");
+        let mut v2 = v1.clone();
+        assert!(v2.append_char('r').is_ok(), "Failed to append 'r'");
+        // Continue with each character in "reasoning>"
+        let mut v3 = v2.clone();
+        for c in "easoning>".chars() {
+            assert!(v3.append_char(c).is_ok(), "Failed to append '{}'", c);
+        }
+        // Try the whole string at once
+        let mut v4 = validator.clone();
+        let result = v4.append("<reasoning>");
+        assert!(result.is_ok(), "Failed to append '<reasoning>': {:?}", result.err());
+    }
+    
+    // Test validation of multiple fragments
+    #[test]
+    fn test_multiple_append() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        // Test sequential appends
+        assert!(validator.append("<").is_ok(), "Failed to append '<'");
+        assert!(validator.append("reasoning").is_ok(), "Failed to append 'reasoning'");
+        assert!(validator.append(">").is_ok(), "Failed to append '>'");
+        assert!(validator.append("test content").is_ok(), "Failed to append content");
+        assert!(validator.append("</reasoning>").is_ok(), "Failed to append closing tag");
+        // Should be at EOF now
+        assert!(validator.eof(), "Should be at EOF after complete XML");
+    }
+    
+    // Test token continuation - critical for the token masking issue
+    #[test]
+    fn test_can_continue_with() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap();
+        // Test at start
+        let validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        assert!(validator.can_continue_with('<'), "Should accept '<' at start");
+        assert!(!validator.can_continue_with('a'), "Should reject 'a' at start");
+        // Test after opening tag
+        let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        validator.append("<reasoning>").unwrap();
+        // Any character should be valid inside a mixed content element
+        assert!(validator.can_continue_with('a'), "Should accept 'a' inside element");
+        assert!(validator.can_continue_with('1'), "Should accept '1' inside element");
+        assert!(validator.can_continue_with('<'), "Should accept '<' inside element");
+    }
+}
+
+#[cfg(test)]
+mod diagnostic_tests {
+    use super::*;
+    use std::sync::Arc;
+
+    // Simple schema with just the reasoning tag for cleaner debugging
+    fn get_simple_schema() -> String {
+        r#"<?xml version="1.0" encoding="UTF-8"?>
+        <xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema" elementFormDefault="qualified">
+          <xs:element name="reasoning">
+            <xs:complexType mixed="true">
+              <xs:sequence>
+                <xs:any minOccurs="0" maxOccurs="unbounded" processContents="skip"/>
+              </xs:sequence>
+            </xs:complexType>
+          </xs:element>
+        </xs:schema>"#.to_string()
+    }
+
+    // Test every character in an opening tag one by one
+    #[test]
+    fn test_char_by_char_open_tag() {
+        let schema_text = get_simple_schema();
+        let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        
+        // Test each character individually
+        assert!(validator.can_continue_with('<'), "Should accept '<'");
+        assert!(validator.append_char('<').is_ok(), "Failed to append '<'");
+        
+        assert!(validator.can_continue_with('r'), "Should accept 'r'");
+        assert!(validator.append_char('r').is_ok(), "Failed to append 'r'");
+        
+        for c in "easoning".chars() {
+            assert!(validator.can_continue_with(c), "Should accept '{}'", c);
+            assert!(validator.append_char(c).is_ok(), "Failed to append '{}'", c);
+        }
+        
+        // This is where many tests are failing - let's check the state before and after
+        println!("State before '>': {:?}", validator.current_tokens);
+        let result = validator.append_char('>');
+        println!("Result of append '>': {:?}", result);
+        if result.is_ok() {
+            println!("State after '>': {:?}", validator.current_tokens);
+        }
+        assert!(result.is_ok(), "Failed to append '>'");
+    }
+
+    // Test the token transitions specifically
+    #[test]
+    fn test_token_transitions() {
+        // Create minimal elements for testing token transitions directly
+        let test_element = Arc::new(schema::XsElement {
+            name: "reasoning".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: Some("true".to_string()),
+                text: None,
+                xs_attribute: None,
+                xs_sequence: Some(schema::XsSequence {
+                    text: None,
+                    xs_any: schema::XsAny {
+                        min_occurs: "0".to_string(),
+                        max_occurs: "unbounded".to_string(),
+                        process_contents: "skip".to_string(),
+                    },
+                }),
+            },
+        });
+        
+        // Test the transition from ElementOpenName to ElementOpenEnd
+        let element_name = token::ElementOpenName::new(Arc::clone(&test_element), "reasoning".to_string()).unwrap();
+        let tokens = element_name.append('>');
+        assert!(!tokens.is_empty(), "Should produce valid tokens after '>'");
+        println!("Token after appending '>' to ElementOpenName: {:?}", tokens[0]);
+        
+        // Check what type the token is
+        match &tokens[0] {
+            Token::ElementOpen(element_open) => {
+                println!("Successfully transitioned to ElementOpen: {:?}", element_open);
+            },
+            other => {
+                panic!("Expected ElementOpen, got {:?}", other);
+            }
+        }
+    }
+
+    // Test element self-closing functionality
+    #[test]
+    fn test_element_self_close_transition() {
+        let test_element = Arc::new(schema::XsElement {
+            name: "stop".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: None,
+                text: None,
+                xs_attribute: None,
+                xs_sequence: None,
+            },
+        });
+        
+        // Test transition from ElementName to ElementSelfClose
+        let element_name = token::ElementOpenName::new(Arc::clone(&test_element), "stop".to_string()).unwrap();
+        let tokens = element_name.append('/');
+        assert!(!tokens.is_empty(), "Should handle '/' after element name");
+        
+        match &tokens[0] {
+            Token::ElementSelfClose(element_self_close) => {
+                println!("Successfully transitioned to ElementSelfClose");
+                // Now test transition to EndOfFile
+                let next_tokens = element_self_close.clone().append('>');
+                assert!(!next_tokens.is_empty(), "Should handle '>' after '/'");
+                
+                match &next_tokens[0] {
+                    Token::EndOfFile(_) => println!("Successfully transitioned to EndOfFile"),
+                    other => panic!("Expected EndOfFile, got {:?}", other),
+                }
+            },
+            other => panic!("Expected ElementSelfClose, got {:?}", other),
+        }
+    }
+
+    // Test whitespace handling
+    #[test]
+    fn test_whitespace_handling() {
+        let schema_text = get_simple_schema();
+        let mut validator = XmlSchemaValidator::new(&schema_text).unwrap();
+        
+        // Test opening tag with whitespace
+        for c in "<reasoning ".chars() {
+            println!("Appending character: '{}'", c);
+            assert!(validator.append_char(c).is_ok(), "Failed to append '{}'", c);
+        }
+        
+        println!("State after '<reasoning ': {:?}", validator.current_tokens);
+        
+        // Test completion with >
+        assert!(validator.can_continue_with('>'), "Should accept '>' after whitespace");
+        assert!(validator.append_char('>').is_ok(), "Failed to append '>' after whitespace");
+    }
+
+    // Test attribute handling
+    #[test]
+    fn test_attribute_handling() {
+        let schema_text = r#"<?xml version="1.0" encoding="UTF-8"?>
+        <xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema" elementFormDefault="qualified">
+          <xs:element name="delete">
+            <xs:complexType>
+              <xs:attribute name="id" type="xs:string" use="required"/>
+            </xs:complexType>
+          </xs:element>
+        </xs:schema>"#;
+        
+        let mut validator = XmlSchemaValidator::new(schema_text).unwrap();
+        
+        // Test opening tag with attribute
+        for c in "<delete ".chars() {
+            assert!(validator.append_char(c).is_ok(), "Failed to append '{}'", c);
+        }
+        
+        // Test attribute name
+        for c in "id".chars() {
+            assert!(validator.can_continue_with(c), "Should accept '{}' in attribute name", c);
+            assert!(validator.append_char(c).is_ok(), "Failed to append '{}' in attribute name", c);
+        }
+        
+        // Test attribute equals and quotes
+        for c in "=\"123\"".chars() {
+            println!("Testing attribute character: '{}'", c);
+            assert!(validator.can_continue_with(c), "Should accept '{}' in attribute value", c);
+            assert!(validator.append_char(c).is_ok(), "Failed to append '{}' in attribute value", c);
+        }
+        
+        // Test closing the self-closing tag
+        println!("Current state before closing: {:?}", validator.current_tokens);
+        for c in "/>".chars() {
+            assert!(validator.can_continue_with(c), "Should accept '{}' to close tag", c);
+            assert!(validator.append_char(c).is_ok(), "Failed to append '{}' to close tag", c);
+        }
+        
+        // Check if we've reached EOF
+        assert!(validator.eof(), "Should be at EOF after self-closing tag");
+    }
+
+    // Test for the specific error patterns seen in the test output
+    #[test]
+    fn test_specific_error_cases() {
+        let schema_text = std::fs::read_to_string("../../action_schema.xsd").unwrap_or_else(|_| get_simple_schema());
+        
+        // Test case 1: Simple opening tag (noticed in `test_element_open`)
+        let input1 = "<reasoning>";
+        let mut validator1 = XmlSchemaValidator::new(&schema_text).unwrap();
+        let result1 = validator1.append(input1);
+        println!("Result appending '<reasoning>': {:?}", result1);
+        assert!(result1.is_ok(), "Failed on '<reasoning>'");
+        
+        // Test case 2: Element with whitespace (noticed in `test_element_open_whitespace`)
+        let input2 = "<reasoning >";
+        let mut validator2 = XmlSchemaValidator::new(&schema_text).unwrap();
+        let result2 = validator2.append(input2);
+        println!("Result appending '<reasoning >': {:?}", result2);
+        assert!(result2.is_ok(), "Failed on '<reasoning >'");
+        
+        // Test case 3: Self-closing tags (noticed in `test_read_stdin_self_closing`)
+        let input3 = "<read_stdin/>";
+        let mut validator3 = XmlSchemaValidator::new(&schema_text).unwrap();
+        let result3 = validator3.append(input3);
+        println!("Result appending '<read_stdin/>': {:?}", result3);
+        assert!(result3.is_ok(), "Failed on '<read_stdin/>'");
+        
+        // Test case 4: Attributes (noticed in `test_single_with_attributes`)
+        let input4 = r#"<single timeout="2.5">"#;
+        let mut validator4 = XmlSchemaValidator::new(&schema_text).unwrap();
+        let result4 = validator4.append(input4);
+        println!("Result appending attribute: {:?}", result4);
+        assert!(result4.is_ok(), "Failed on attribute handling");
+    }
+
+    // Test the basic token state machine with a minimal element
+    #[test]
+    fn test_basic_token_operations() {
+        // Create a minimal element for testing
+        let element = Arc::new(schema::XsElement {
+            name: "test".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: Some("true".to_string()),
+                text: None,
+                xs_attribute: None,
+                xs_sequence: None,
+            },
+        });
+        
+        // Create and test each token type
+        let start_of_file = token::StartOfFile::new(Arc::clone(&element));
+        let tokens = start_of_file.append('<');
+        assert!(!tokens.is_empty(), "StartOfFile should accept '<'");
+        
+        let element_start = match &tokens[0] {
+            Token::ElementStart(es) => es.clone(),
+            _ => panic!("Expected ElementStart"),
+        };
+        
+        let tokens = element_start.append('t');
+        assert!(!tokens.is_empty(), "ElementStart should accept 't'");
+        
+        let element_name = match &tokens[0] {
+            Token::ElementName(en) => en.clone(),
+            _ => panic!("Expected ElementName"),
+        };
+        
+        // This is often where things break
+        println!("About to append 'e' to ElementName...");
+        let tokens = element_name.append('e');
+        assert!(!tokens.is_empty(), "ElementName should accept 'e'");
+        
+        // Build the rest of the name
+        let mut current_token = tokens[0].clone();
+        for c in "st".chars() {
+            let next_tokens = current_token.append(c);
+            assert!(!next_tokens.is_empty(), "Token should accept '{}'", c);
+            current_token = next_tokens[0].clone();
+        }
+        
+        // Now try to append '>'
+        println!("About to append '>' to token: {:?}", current_token);
+        let tokens = current_token.append('>');
+        println!("Result of appending '>': {:?}", tokens);
+        assert!(!tokens.is_empty(), "Token should accept '>'");
+        
+        // Verify transition to ElementOpen
+        match &tokens[0] {
+            Token::ElementOpen(_) => println!("Successfully transitioned to ElementOpen"),
+            other => panic!("Expected ElementOpen, got {:?}", other),
+        }
+    }
+    
+    // This test will help identify issues with token creation for specific elements
+    #[test]
+    fn test_element_specific_issues() {
+        if let Ok(schema_text) = std::fs::read_to_string("../../action_schema.xsd") {
+            let validator = XmlSchemaValidator::new(&schema_text).unwrap();
+            
+            // Get all element names in the schema
+            let element_names = validator.get_element_names();
+            println!("Elements in schema: {:?}", element_names);
+            
+            // Test each element separately
+            for name in &element_names {
+                println!("\nTesting element: {}", name);
+                
+                let mut element_validator = XmlSchemaValidator::new(&schema_text).unwrap();
+                let open_tag = format!("<{}>", name);
+                
+                // Test character by character
+                println!("Testing character by character for <{}>", name);
+                let mut progress = String::new();
+                for c in open_tag.chars() {
+                    progress.push(c);
+                    let result = element_validator.can_continue_with(c);
+                    println!("  Can continue with '{}' after '{}': {}", 
+                              c, &progress[..progress.len()-1], result);
+                    
+                    if !result {
+                        // This is helpful to diagnose where the process fails
+                        println!("  ❌ Failed on character '{}' for element '{}'", c, name);
+                        break;
+                    }
+                    
+                    if let Err(e) = element_validator.append_char(c) {
+                        println!("  ❌ Failed to append '{}' for element '{}': {:?}", c, name, e);
+                        break;
+                    }
+                }
+                
+                // Test complete tag at once
+                let mut complete_validator = XmlSchemaValidator::new(&schema_text).unwrap();
+                match complete_validator.append(&open_tag) {
+                    Ok(_) => println!("  ✅ Successfully appended <{}>", name),
+                    Err(e) => println!("  ❌ Failed to append <{}>: {:?}", name, e),
+                }
+            }
+        } else {
+            println!("Schema file not found, skipping element-specific tests");
+        }
+    }
 }

lib/xml_schema_validator/src/py_xml_logits_processor_core.rs

@@ -23,10 +23,13 @@ impl XmlLogitsProcessorCore {
         })
     }
     
+    /// Append a fragment of XML to the validator
+    /// Returns true if the append was successful, false otherwise
     fn append(&mut self, fragment: &str) -> bool {
         self.xml_schema_validator.append(fragment).is_ok()
     }
 
+    /// Get a list of tokens that would lead to invalid XML
     fn get_invalid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
         Ok(self.tokens.iter()
             .filter_map(|token| {
@@ -38,21 +41,412 @@ impl XmlLogitsProcessorCore {
             })
             .collect::<Vec<_>>())
     }
+    
+    /// Get a list of tokens that would lead to valid XML
+    fn get_valid_tokens(&mut self, py: Python<'_>) -> PyResult<Vec<PyObject>> {
+        Ok(self.tokens.iter()
+            .filter_map(|token| {
+                if self.xml_schema_validator.clone().append(&token.1).is_ok() {
+                    Some(token.0.clone_ref(py))
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>())
+    }
 
+    /// Check if the validator has reached the end of the XML
     fn eof(&self) -> PyResult<bool> {
         Ok(self.xml_schema_validator.eof())
     }
 
+    /// Create a copy of the processor
     fn copy(&self, py: Python<'_>) -> PyResult<Self> {
         Ok(Self {
             xml_schema_validator: self.xml_schema_validator.clone(),
             tokens: self.tokens.iter().map(|(a, b)|(a.clone_ref(py), b.clone())).collect()
         })
     }
+    
+    /// Get a list of all element names in the schema
+    fn get_element_names(&self) -> PyResult<Vec<String>> {
+        Ok(self.xml_schema_validator.get_element_names())
+    }
+    
+    /// Check if a specific character is valid as the next character
+    fn can_continue_with(&self, c: &str) -> PyResult<bool> {
+        if let Some(first_char) = c.chars().next() {
+            Ok(self.xml_schema_validator.can_continue_with(first_char))
+        } else {
+            Ok(false)
+        }
+    }
+    
+    /// Reset the validator to its initial state
+    fn reset(&mut self) -> PyResult<()> {
+        // Create a new validator with the same schema elements
+        let schema_text = r#"<?xml version="1.0" encoding="UTF-8"?>
+<xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema" elementFormDefault="qualified">
+  <xs:element name="root">
+    <xs:complexType mixed="true">
+      <xs:sequence>
+        <xs:any minOccurs="0" maxOccurs="unbounded" processContents="skip"/>
+      </xs:sequence>
+    </xs:complexType>
+  </xs:element>
+</xs:schema>"#;
+        
+        match crate::XmlSchemaValidator::new(schema_text) {
+            Ok(validator) => {
+                self.xml_schema_validator = validator;
+                Ok(())
+            },
+            Err(e) => Err(pyo3::exceptions::PyRuntimeError::new_err(format!(
+                "Failed to reset validator: {}", e
+            ))),
+        }
+    }
 }
 
 /// Register the class with the Python module
 pub fn register_module(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<XmlLogitsProcessorCore>()?;
     Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use pyo3::Python;
+    use pyo3::types::PyDict;
+
+    fn create_test_processor<'py>(py: Python<'py>) -> PyResult<XmlLogitsProcessorCore> {
+        let tokens = PyDict::new(py);
+        // Add some basic tokens
+        tokens.set_item(0, "<")?;
+        tokens.set_item(1, "reasoning")?;
+        tokens.set_item(2, ">")?;
+        tokens.set_item(3, " ")?;
+        tokens.set_item(4, "a")?;
+        tokens.set_item(5, "</reasoning>")?;
+        
+        // Simple schema with only reasoning element
+        let schema_text = r#"<?xml version="1.0" encoding="UTF-8"?>
+        <xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema" elementFormDefault="qualified">
+          <xs:element name="reasoning">
+            <xs:complexType mixed="true">
+              <xs:sequence>
+                <xs:any minOccurs="0" maxOccurs="unbounded" processContents="skip"/>
+              </xs:sequence>
+            </xs:complexType>
+          </xs:element>
+        </xs:schema>"#;
+        
+        XmlLogitsProcessorCore::new(tokens.into(), schema_text)
+    }
+
+    #[test]
+    fn test_append_valid_xml() {
+        Python::with_gil(|py| {
+            let mut processor = create_test_processor(py).unwrap();
+            
+            // Test various valid XML fragments
+            assert!(processor.append("<reasoning>"), "Should accept opening tag");
+            
+            let mut processor = create_test_processor(py).unwrap();
+            assert!(processor.append("<reasoning>test</reasoning>"), 
+                   "Should accept complete element with content");
+            
+            // Test with real schema from disk if available
+            if let Ok(schema_text) = std::fs::read_to_string("../../action_schema.xsd") {
+                let tokens = PyDict::new(py);
+                tokens.set_item(0, "<").unwrap();
+                tokens.set_item(1, "reasoning").unwrap();
+                tokens.set_item(2, ">").unwrap();
+                
+                let mut processor = XmlLogitsProcessorCore::new(tokens.into(), &schema_text).unwrap();
+                assert!(processor.append("<reasoning>"), 
+                       "Should accept <reasoning> with actual schema");
+            }
+        });
+    }
+
+    #[test]
+    fn test_append_invalid_xml() {
+        Python::with_gil(|py| {
+            let mut processor = create_test_processor(py).unwrap();
+            
+            // Test invalid XML fragments
+            assert!(!processor.append("<invalid>"), "Should reject invalid element");
+            
+            let mut processor = create_test_processor(py).unwrap();
+            assert!(!processor.append("<reasoning></invalid>"), 
+                   "Should reject mismatched tags");
+        });
+    }
+
+    #[test]
+    fn test_get_valid_tokens() {
+        Python::with_gil(|py| {
+            let mut processor = create_test_processor(py).unwrap();
+            
+            // At start, only "<" is valid
+            let valid_tokens = processor.get_valid_tokens(py).unwrap();
+            
+            let contains_0 = valid_tokens.iter().any(|obj| {
+                let obj_id: Option<i64> = obj.extract(py).ok();
+                obj_id == Some(0)
+            });
+            
+            assert!(contains_0, "Token '<' should be valid at start");
+            assert_eq!(valid_tokens.len(), 1, "Only one token should be valid at start");
+            
+            // After <reasoning>, any character should be valid in mixed content
+            let mut processor = create_test_processor(py).unwrap();
+            assert!(processor.append("<reasoning>"), "Should be able to append <reasoning>");
+            
+            let valid_tokens = processor.get_valid_tokens(py).unwrap();
+            
+            let contains_3 = valid_tokens.iter().any(|obj| {
+                let obj_id: Option<i64> = obj.extract(py).ok();
+                obj_id == Some(3)
+            });
+            
+            let contains_4 = valid_tokens.iter().any(|obj| {
+                let obj_id: Option<i64> = obj.extract(py).ok();
+                obj_id == Some(4)
+            });
+            
+            assert!(contains_3, "Token ' ' should be valid inside mixed content");
+            assert!(contains_4, "Token 'a' should be valid inside mixed content");
+        });
+    }
+
+    #[test]
+    fn test_eof_detection() {
+        Python::with_gil(|py| {
+            let processor = create_test_processor(py).unwrap();
+            
+            // At start, not at EOF
+            assert!(!processor.eof().unwrap(), "Should not be at EOF at start");
+            
+            // After complete document, should be at EOF
+            let mut processor = create_test_processor(py).unwrap();
+            // Use a properly formed XML document that matches the schema
+            assert!(processor.append("<reasoning>test</reasoning>"), 
+                "Should accept complete document");
+            
+            // DEBUG: Print the token state to see why EOF detection fails
+            println!("Current tokens: {:?}", processor.xml_schema_validator.current_tokens);
+            
+            assert!(processor.eof().unwrap(), "Should be at EOF after complete document");
+            
+            // After opening tag, should not be at EOF
+            let mut processor = create_test_processor(py).unwrap();
+            assert!(processor.append("<reasoning>"), "Should accept opening tag");
+            assert!(!processor.eof().unwrap(), "Should not be at EOF after opening tag");
+        });
+    }
+
+    #[test]
+    fn test_can_continue_with() {
+        Python::with_gil(|py| {
+            let processor = create_test_processor(py).unwrap();
+            
+            // At start, only '<' is valid
+            assert!(processor.can_continue_with("<").unwrap(), "Should accept '<' at start");
+            assert!(!processor.can_continue_with("a").unwrap(), "Should reject 'a' at start");
+            
+            // Test after opening tag
+            let mut processor = create_test_processor(py).unwrap();
+            assert!(processor.append("<reasoning>"), "Should accept opening tag");
+            
+            // In mixed content, any character should be valid
+            assert!(processor.can_continue_with("a").unwrap(), 
+                   "Should accept 'a' in mixed content");
+            assert!(processor.can_continue_with("<").unwrap(), 
+                   "Should accept '<' in mixed content");
+        });
+    }
+
+    #[test]
+    fn test_copy() {
+        Python::with_gil(|py| {
+            let processor = create_test_processor(py).unwrap();
+            
+            // Make a copy
+            let copy = processor.copy(py).unwrap();
+            
+            // Original and copy should behave the same
+            assert_eq!(processor.eof().unwrap(), copy.eof().unwrap(), 
+                      "Copy should have same EOF state");
+            
+            // Modify copy, original should be unchanged
+            let mut copy = processor.copy(py).unwrap();
+            assert!(copy.append("<reasoning>"), "Copy should accept XML");
+            assert!(!processor.eof().unwrap(), "Original should be unchanged");
+            assert!(!copy.eof().unwrap(), "Copy should be modified");
+        });
+    }
+
+    #[test]
+    fn test_get_element_names() {
+        Python::with_gil(|py| {
+            let processor = create_test_processor(py).unwrap();
+            
+            let names = processor.get_element_names().unwrap();
+            assert!(names.contains(&"reasoning".to_string()), 
+                   "Should contain 'reasoning' element");
+            assert_eq!(names.len(), 1, "Should have exactly one element");
+            
+            // Test with real schema from disk if available
+            if let Ok(schema_text) = std::fs::read_to_string("../../action_schema.xsd") {
+                let tokens = PyDict::new(py);
+                tokens.set_item(0, "<").unwrap();
+                
+                let processor = XmlLogitsProcessorCore::new(tokens.into(), &schema_text).unwrap();
+                let names = processor.get_element_names().unwrap();
+                
+                // Print out all elements from schema for debugging
+                println!("Elements in schema: {:?}", names);
+                
+                // Check if reasoning exists
+                assert!(names.contains(&"reasoning".to_string()), 
+                       "Real schema should contain 'reasoning'");
+            }
+        });
+    }
+    
+    #[test]
+    fn test_debug_reasoning_tag_validation() {
+        Python::with_gil(|py| {
+            // Create a simple schema with just the reasoning element for testing
+            let schema_text = r#"<?xml version="1.0" encoding="UTF-8"?>
+            <xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema" elementFormDefault="qualified">
+              <xs:element name="reasoning">
+                <xs:complexType mixed="true">
+                  <xs:sequence>
+                    <xs:any minOccurs="0" maxOccurs="unbounded" processContents="skip"/>
+                  </xs:sequence>
+                </xs:complexType>
+              </xs:element>
+            </xs:schema>"#;
+            
+            // Create a minimal token dictionary
+            let tokens = PyDict::new(py);
+            tokens.set_item(0, "<").unwrap();
+            tokens.set_item(1, "r").unwrap();
+            tokens.set_item(2, "e").unwrap();
+            tokens.set_item(3, "a").unwrap();
+            tokens.set_item(4, "s").unwrap();
+            tokens.set_item(5, "o").unwrap();
+            tokens.set_item(6, "n").unwrap();
+            tokens.set_item(7, "i").unwrap();
+            tokens.set_item(8, "n").unwrap();
+            tokens.set_item(9, "g").unwrap();
+            tokens.set_item(10, ">").unwrap();
+            
+            // Create processor with this schema and tokens
+            let processor = XmlLogitsProcessorCore::new(tokens.clone().into(), schema_text).unwrap();
+            
+            // Test basic operations to verify core functionality works
+            let element_names = processor.get_element_names().unwrap();
+            println!("Element names in test schema: {:?}", element_names);
+            assert!(element_names.contains(&"reasoning".to_string()), 
+                   "Test schema should contain 'reasoning'");
+            
+            // Now try to append each character of "<reasoning>" one by one
+            let chars = vec!['<', 'r', 'e', 'a', 's', 'o', 'n', 'i', 'n', 'g', '>'];
+            let mut full_text = String::new();
+            let processor = processor.copy(py).unwrap();
+            
+            for c in chars {
+                full_text.push(c);
+                let fragment = c.to_string();
+                let result = processor.can_continue_with(&fragment).unwrap();
+                println!("Can continue with '{}' after '{}': {}", 
+                        c, &full_text[..full_text.len()-1], result);
+                
+                // Create a new processor to test appending the fragment built so far
+                let mut fresh_processor = XmlLogitsProcessorCore::new(tokens.clone().into(), schema_text).unwrap();
+                let append_result = fresh_processor.append(&full_text);
+                println!("Appending full text '{}': {}", full_text, append_result);
+            }
+            
+            // Now test with actual schema from disk if available
+            if let Ok(actual_schema) = std::fs::read_to_string("../../action_schema.xsd") {
+                println!("\nTesting with actual schema from disk:");
+                
+                let processor = XmlLogitsProcessorCore::new(tokens.clone().into(), &actual_schema).unwrap();
+                let element_names = processor.get_element_names().unwrap();
+                println!("Element names in actual schema: {:?}", element_names);
+                
+                if element_names.contains(&"reasoning".to_string()) {
+                    println!("'reasoning' element found in actual schema");
+                    
+                    // Test appending "<reasoning>" to a fresh processor
+                    let mut processor = processor.copy(py).unwrap();
+                    let result = processor.append("<reasoning>");
+                    println!("Appending '<reasoning>' to fresh processor: {}", result);
+                    
+                    // Test character by character
+                    let mut text = String::new();
+                    for c in "<reasoning>".chars() {
+                        text.push(c);
+                        let result = XmlLogitsProcessorCore::new(tokens.clone().into(), &actual_schema).unwrap()
+                                      .append(&text);
+                        println!("Appending '{}' to fresh processor: {}", text, result);
+                    }
+                } else {
+                    println!("Warning: 'reasoning' element NOT found in actual schema!");
+                }
+            }
+        });
+    }
+
+    #[test]
+    fn test_debug_token_processing() {
+        Python::with_gil(|py| {
+            if let Ok(schema_text) = std::fs::read_to_string("../../action_schema.xsd") {
+                // Create a dictionary mapping token IDs to their string values
+                // Simulating how the tokenizer vocabulary works
+                let tokens = PyDict::new(py);
+                tokens.set_item(1000, "<").unwrap();
+                tokens.set_item(1001, "reasoning").unwrap();
+                tokens.set_item(1002, ">").unwrap();
+                tokens.set_item(1003, " ").unwrap();
+                tokens.set_item(1004, "test").unwrap();
+                tokens.set_item(1005, "</").unwrap();
+                tokens.set_item(1006, "</reasoning>").unwrap();
+                
+                let mut processor = XmlLogitsProcessorCore::new(tokens.clone().into(), &schema_text).unwrap();
+                
+                // Test which tokens are valid at the start
+                let invalid_tokens = processor.get_invalid_tokens(py).unwrap();
+                let valid_tokens = processor.get_valid_tokens(py).unwrap();
+                
+                println!("At start - valid tokens count: {}", valid_tokens.len());
+                println!("At start - invalid tokens count: {}", invalid_tokens.len());
+                
+                // Create a new processor and progress to after "<reasoning>"
+                let mut processor = XmlLogitsProcessorCore::new(tokens.into(), &schema_text).unwrap();
+                let result = processor.append("<reasoning>");
+                println!("Append '<reasoning>': {}", result);
+                
+                if result {
+                    // Now get valid/invalid tokens after "<reasoning>"
+                    let invalid_tokens = processor.get_invalid_tokens(py).unwrap();
+                    let valid_tokens = processor.get_valid_tokens(py).unwrap();
+                    
+                    println!("After '<reasoning>' - valid tokens count: {}", valid_tokens.len());
+                    println!("After '<reasoning>' - invalid tokens count: {}", invalid_tokens.len());
+                    
+                    // At this point we'd expect to have more valid tokens than at the start
+                    // since mixed content allows any text
+                }
+            } else {
+                println!("Warning: Couldn't find schema file for testing");
+            }
+        });
+    }
 }

lib/xml_schema_validator/src/py_xml_schema_validator.rs

@@ -17,6 +17,10 @@ impl XmlSchemaValidator {
         Ok(Self { validator })
     }
 
+    /// Append a fragment of XML to the validator
+    /// Returns a tuple containing:
+    /// - A boolean indicating success
+    /// - An optional error message
     fn append(&mut self, fragment: &str) -> PyResult<(bool, Option<String>)> {
         match self.validator.append(fragment) {
             Ok(()) => PyResult::Ok((true, None)),
@@ -24,15 +28,40 @@ impl XmlSchemaValidator {
         }
     }
 
+    /// Check if the validator has reached the end of the XML
     fn eof(&self) -> PyResult<bool> {
         Ok(self.validator.eof())
     }
 
+    /// Create a deep copy of this validator
     fn copy(&self) -> PyResult<Self> {
         Ok(Self {
             validator: self.validator.clone(),
         })
     }
+    
+    /// Get a list of all element names in the schema
+    fn get_element_names(&self) -> PyResult<Vec<String>> {
+        Ok(self.validator.get_element_names())
+    }
+    
+    /// Check if a specific character is valid as the next character
+    fn can_continue_with(&self, c: &str) -> PyResult<bool> {
+        if let Some(first_char) = c.chars().next() {
+            Ok(self.validator.can_continue_with(first_char))
+        } else {
+            Ok(false)
+        }
+    }
+    
+    /// Validate an entire XML string and check if it's valid according to the schema
+    fn validate(&self, xml_string: &str) -> PyResult<bool> {
+        let mut validator_clone = self.validator.clone();
+        match validator_clone.append(xml_string) {
+            Ok(()) => Ok(validator_clone.eof()),  // Valid only if we reached EOF
+            Err(_) => Ok(false),  // Invalid XML
+        }
+    }
 }
 
 /// Register the class with the Python module

lib/xml_schema_validator/src/token/attribute_name.rs

@@ -0,0 +1,150 @@
+use std::sync::Arc;
+use crate::*;
+
+/// Represents an attribute name of an XML element
+#[derive(Clone, Debug)]
+pub struct AttributeName {
+    /// The element that owns this attribute
+    element: Arc<schema::XsElement>,
+    /// Buffer containing the attribute name characters processed so far
+    buffer: String,
+}
+
+impl AttributeName {
+    /// Create a new AttributeName with an empty buffer
+    pub fn new(element: Arc<schema::XsElement>) -> Self {
+        Self {
+            element,
+            buffer: String::new(),
+        }
+    }
+    
+    /// Create a new AttributeName with the given buffer
+    pub fn with_buffer(element: Arc<schema::XsElement>, buffer: String) -> Self {
+        Self {
+            element,
+            buffer,
+        }
+    }
+
+    /// Append a character to the attribute name and return possible continuations
+    pub fn append(self, c: char) -> Vec<Token> {
+        if c.is_whitespace() && self.buffer.is_empty() {
+            return vec![Token::AttributeName(self)];
+        } else if c == '=' {
+            // Found the equals sign, transition to attribute value
+            // We need to verify that this attribute name is valid for this element
+            if let Some(attributes) = &self.element.xs_complex_type.xs_attribute {
+                for attr in attributes {
+                    if attr.name == self.buffer {
+                        // Valid attribute, transition to attribute value
+                        return vec![Token::AttributeValue(AttributeValue::new(
+                            Arc::clone(&self.element),
+                            self.buffer,
+                        ))];
+                    }
+                }
+            }
+            // Invalid attribute name
+            vec![]
+        } else {
+            // Continue building the attribute name
+            let new_buffer = self.buffer + &c.to_string();
+            
+            // Check if this could be a valid attribute for this element
+            if let Some(attributes) = &self.element.xs_complex_type.xs_attribute {
+                for attr in attributes {
+                    if attr.name.starts_with(&new_buffer) {
+                        // This could be a valid attribute, continue parsing
+                        return vec![Token::AttributeName(AttributeName::with_buffer(
+                            Arc::clone(&self.element),
+                            new_buffer,
+                        ))];
+                    }
+                }
+            }
+            
+            // No matching attribute found
+            vec![]
+        }
+    }
+}
+
+#[cfg(test)]
+mod attribute_name_tests {
+    use super::*;
+    use std::sync::Arc;
+    use crate::schema;
+
+    fn create_element_with_attributes() -> Arc<schema::XsElement> {
+        let element = schema::XsElement {
+            name: "delete".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: None,
+                text: None,
+                xs_attribute: Some(vec![
+                    schema::XsAttribute {
+                        name: "id".to_string(),
+                        xs_attribute_type: "xs:string".to_string(),
+                        xs_attribute_use: "required".to_string(),
+                    }
+                ]),
+                xs_sequence: None,
+            },
+        };
+        Arc::new(element)
+    }
+
+    #[test]
+    fn test_attribute_name_append() {
+        let element = create_element_with_attributes();
+        let token = AttributeName::new(Arc::clone(&element));
+        
+        // Build attribute name character by character
+        let mut current_token = token;
+        for c in "id".chars() {
+            let next_tokens = current_token.append(c);
+            assert!(!next_tokens.is_empty(), "Should accept attribute name character '{}'", c);
+            
+            match &next_tokens[0] {
+                Token::AttributeName(next) => current_token = next.clone(),
+                _ => panic!("Expected AttributeName, got {:?}", next_tokens[0]),
+            }
+        }
+        
+        // Test equals sign
+        let next_tokens = current_token.append('=');
+        assert!(!next_tokens.is_empty(), "Should accept '=' after attribute name");
+        
+        match &next_tokens[0] {
+            Token::AttributeValue(_) => (),
+            _ => panic!("Expected AttributeValue, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_attribute_name_whitespace() {
+        let element = create_element_with_attributes();
+        let token = AttributeName::new(Arc::clone(&element));
+        
+        // Test with whitespace
+        let next_tokens = token.append(' ');
+        assert!(!next_tokens.is_empty(), "Should accept whitespace in attribute name");
+        
+        match &next_tokens[0] {
+            Token::AttributeName(_) => (),
+            _ => panic!("Expected AttributeName, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_invalid_attribute_name() {
+        let element = create_element_with_attributes();
+        let token = AttributeName::new(Arc::clone(&element));
+        
+        // Test with invalid attribute name
+        let next_tokens = token.append('x'); // 'x' doesn't start any valid attribute
+        assert!(next_tokens.is_empty(), "Should reject invalid attribute name");
+    }
+}

lib/xml_schema_validator/src/token/attribute_value.rs

@@ -0,0 +1,215 @@
+use std::sync::Arc;
+use crate::*;
+
+/// Represents an attribute value of an XML element
+#[derive(Clone, Debug)]
+pub struct AttributeValue {
+    /// The element that owns this attribute
+    element: Arc<schema::XsElement>,
+    /// The name of the attribute being processed
+    attribute_name: String,
+    /// Buffer containing the attribute value characters processed so far
+    buffer: String,
+    /// Whether we've encountered the opening quote
+    in_quotes: bool,
+    /// Whether we've seen the closing quote
+    closed: bool,
+}
+
+impl AttributeValue {
+    /// Create a new AttributeValue
+    pub fn new(element: Arc<schema::XsElement>, attribute_name: String) -> Self {
+        Self {
+            element,
+            attribute_name,
+            buffer: String::new(),
+            in_quotes: false,
+            closed: false,
+        }
+    }
+
+    /// Append a character to the attribute value and return possible continuations
+    pub fn append(self, c: char) -> Vec<Token> {
+        if self.closed {
+            // We've already closed the attribute value with a quote
+            if c.is_whitespace() {
+                // Check if there are more attributes to parse
+                if let Some(attributes) = &self.element.xs_complex_type.xs_attribute {
+                    // If there are other required attributes that haven't been set yet
+                    for attr in attributes {
+                        if attr.name != self.attribute_name && attr.xs_attribute_use == "required" {
+                            return vec![Token::AttributeName(AttributeName::new(Arc::clone(&self.element)))];
+                        }
+                    }
+                }
+                
+                // No more required attributes, whitespace after attribute value, continue to next state
+                return vec![
+                    Token::ElementOpen(ElementOpenEnd::new(Arc::clone(&self.element))),
+                    Token::ElementSelfClose(ElementSelfClose::new())
+                ];
+            } else if c == '>' {
+                // End of opening tag
+                return vec![Token::ElementOpen(ElementOpenEnd::new(Arc::clone(&self.element)))];
+            } else if c == '/' {
+                // Beginning of self-closing tag
+                return vec![Token::ElementSelfClose(ElementSelfClose::new())];
+            } else {
+                // Unexpected character after attribute value
+                return vec![];
+            }
+        }
+        
+        if !self.in_quotes {
+            if c.is_whitespace() {
+                // Skip whitespace before the quotes
+                return vec![Token::AttributeValue(Self {
+                    element: Arc::clone(&self.element),
+                    attribute_name: self.attribute_name,
+                    buffer: self.buffer,
+                    in_quotes: false,
+                    closed: false,
+                })];
+            } else if c == '"' || c == '\'' {
+                // Start of quotes
+                return vec![Token::AttributeValue(Self {
+                    element: Arc::clone(&self.element),
+                    attribute_name: self.attribute_name,
+                    buffer: self.buffer,
+                    in_quotes: true,
+                    closed: false,
+                })];
+            } else {
+                // Unexpected character before quotes
+                return vec![];
+            }
+        } else {
+            // We're inside quotes
+            if c == '"' || c == '\'' {
+                // End of quotes, validate the attribute value
+                // In a real implementation, we'd check if the value is valid for this attribute type
+                return vec![Token::AttributeValue(Self {
+                    element: Arc::clone(&self.element),
+                    attribute_name: self.attribute_name,
+                    buffer: self.buffer,
+                    in_quotes: true,
+                    closed: true,
+                })];
+            } else {
+                // Continue building the attribute value
+                let new_buffer = self.buffer + &c.to_string();
+                return vec![Token::AttributeValue(Self {
+                    element: Arc::clone(&self.element),
+                    attribute_name: self.attribute_name,
+                    buffer: new_buffer,
+                    in_quotes: true,
+                    closed: false,
+                })];
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod attribute_value_tests {
+    use super::*;
+    use std::sync::Arc;
+    use crate::schema;
+
+    fn create_element_with_attributes() -> Arc<schema::XsElement> {
+        let element = schema::XsElement {
+            name: "delete".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: None,
+                text: None,
+                xs_attribute: Some(vec![
+                    schema::XsAttribute {
+                        name: "id".to_string(),
+                        xs_attribute_type: "xs:string".to_string(),
+                        xs_attribute_use: "required".to_string(),
+                    }
+                ]),
+                xs_sequence: None,
+            },
+        };
+        Arc::new(element)
+    }
+
+    #[test]
+    fn test_attribute_value_quotes() {
+        let element = create_element_with_attributes();
+        let token = AttributeValue::new(Arc::clone(&element), "id".to_string());
+        
+        // Test with opening quote
+        let next_tokens = token.append('"');
+        assert!(!next_tokens.is_empty(), "Should accept opening quote");
+        
+        let token = next_tokens[0].clone();
+        match token.clone() {
+            Token::AttributeValue(value) => {
+                assert!(value.in_quotes, "Should mark as in quotes");
+                assert!(!value.closed, "Should not be closed yet");
+            },
+            _ => panic!("Expected AttributeValue, got {:?}", token),
+        }
+        
+        // Test with content inside quotes
+        let next_tokens = token.append('a');
+        assert!(!next_tokens.is_empty(), "Should accept content inside quotes");
+        
+        let token = next_tokens[0].clone();
+        
+        // Test with closing quote
+        let next_tokens = token.append('"');
+        assert!(!next_tokens.is_empty(), "Should accept closing quote");
+        
+        let token = next_tokens[0].clone();
+        match token {
+            Token::AttributeValue(value) => {
+                assert!(value.closed, "Should be closed");
+            },
+            _ => panic!("Expected AttributeValue, got {:?}", token),
+        }
+    }
+
+    #[test]
+    fn test_after_attribute_value_closed() {
+        let element = create_element_with_attributes();
+        let mut token = AttributeValue::new(Arc::clone(&element), "id".to_string());
+        token.in_quotes = true;
+        token.closed = true;
+        
+        // Test with space after closed attribute value
+        let next_tokens = token.append(' ');
+        assert!(!next_tokens.is_empty(), "Should accept space after closed attribute");
+        
+        // Test with > after closed attribute value
+        let token = AttributeValue::new(Arc::clone(&element), "id".to_string());
+        let mut token = token;
+        token.in_quotes = true;
+        token.closed = true;
+        
+        let next_tokens = token.append('>');
+        assert!(!next_tokens.is_empty(), "Should accept '>' after closed attribute");
+        
+        match &next_tokens[0] {
+            Token::ElementOpen(_) => (),
+            _ => panic!("Expected ElementOpen, got {:?}", next_tokens[0]),
+        }
+        
+        // Test with / after closed attribute value (for self-closing)
+        let token = AttributeValue::new(Arc::clone(&element), "id".to_string());
+        let mut token = token;
+        token.in_quotes = true;
+        token.closed = true;
+        
+        let next_tokens = token.append('/');
+        assert!(!next_tokens.is_empty(), "Should accept '/' after closed attribute");
+        
+        match &next_tokens[0] {
+            Token::ElementSelfClose(_) => (),
+            _ => panic!("Expected ElementSelfClose, got {:?}", next_tokens[0]),
+        }
+    }
+}

lib/xml_schema_validator/src/token/element_close_end.rs

@@ -1,21 +0,0 @@
-use crate::*;
-
-/// Represents the end of a closing element tag
-#[derive(Clone, Debug)]
-pub struct ElementCloseEnd {
-}
-
-impl ElementCloseEnd {
-    pub fn new() -> Self {
-        Self {
-        }
-    }
-
-    pub fn append(self, c: char) -> Vec<Token> {
-        if '>' == c {
-            vec![Token::EndOfFile(EndOfFile::new())]
-        } else {
-            vec![]
-        }
-    }
-}

lib/xml_schema_validator/src/token/element_close_name.rs

@@ -21,32 +21,117 @@ impl ElementCloseName {
     }
 
     pub fn append(self, c: char) -> Vec<Token> {
+        if self.element.name == self.buffer && c.is_whitespace() {
+            return vec![Token::ElementCloseName(self)];
+        }
+        if self.element.name == self.buffer && c == '>' {
+            return vec![Token::EndOfFile(EndOfFile::new())];
+        }
         let new_buffer = self.buffer + &c.to_string();
-        
-        // If still building the element name
         if self.element.name.starts_with(&new_buffer) {
-            // If we've matched the full name
-            if self.element.name == new_buffer {
-                // Now we're expecting '>' to close the tag
-                if c.is_whitespace() {
-                    // Handle whitespace after the element name
-                    vec![Token::Whitespace(Whitespace::new(
-                        vec![Token::ElementCloseEnd(ElementCloseEnd::new())]
-                    ))]
-                } else {
-                    // Directly transition to the closed state
-                    vec![Token::ElementCloseEnd(ElementCloseEnd::new())]
-                }
-            } else {
-                // Still building the element name
-                vec![Token::ElementCloseName(ElementCloseName::new(
-                    Arc::clone(&self.element),
-                    new_buffer,
-                ))]
-            }
+            return vec![Token::ElementCloseName(ElementCloseName::new(
+                Arc::clone(&self.element),
+                new_buffer,
+            ))];
         } else {
-            // The character doesn't match what we expect for this element name
             vec![]
         }
     }
+}
+
+#[cfg(test)]
+mod element_close_name_tests {
+    use super::*;
+    use std::sync::Arc;
+    use crate::schema;
+
+    fn create_test_element() -> Arc<schema::XsElement> {
+        let element = schema::XsElement {
+            name: "reasoning".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: Some("true".to_string()),
+                text: None,
+                xs_attribute: None,
+                xs_sequence: Some(schema::XsSequence {
+                    text: None,
+                    xs_any: schema::XsAny {
+                        min_occurs: "0".to_string(),
+                        max_occurs: "unbounded".to_string(),
+                        process_contents: "skip".to_string(),
+                    },
+                }),
+            },
+        };
+        Arc::new(element)
+    }
+
+    #[test]
+    fn test_close_name_partial() {
+        let element = create_test_element();
+        let token = ElementCloseName::new(Arc::clone(&element), "r".to_string());
+        
+        // Test with next character in name
+        let next_tokens = token.append('e');
+        assert!(!next_tokens.is_empty(), "Should accept next character in name");
+        
+        match &next_tokens[0] {
+            Token::ElementCloseName(name) => {
+                assert_eq!(name.buffer, "re", "Buffer should contain 're'");
+            },
+            _ => panic!("Expected ElementCloseName, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_close_name_complete() {
+        let element = create_test_element();
+        let token = ElementCloseName::new(Arc::clone(&element), "reasonin".to_string());
+        
+        // Test with last character of name
+        let next_tokens = token.append('g');
+        assert!(!next_tokens.is_empty(), "Should accept last character of name");
+        
+        match &next_tokens[0] {
+            Token::ElementCloseName(name) => {
+                assert_eq!(name.buffer, "reasoning", "Buffer should contain 'reasoning'");
+            },
+            _ => panic!("Expected ElementCloseName, got {:?}", next_tokens[0]),
+        }
+
+        // Test with > after complete name
+        let token = ElementCloseName::new(Arc::clone(&element), "reasoning".to_string());
+        let next_tokens = token.append('>');
+        assert!(!next_tokens.is_empty(), "Should accept '>' after complete name");
+        
+        match &next_tokens[0] {
+            Token::EndOfFile(_) => (),
+            _ => panic!("Expected EndOfFile, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_close_name_whitespace() {
+        let element = create_test_element();
+        let token = ElementCloseName::new(Arc::clone(&element), "reasoning".to_string());
+        
+        // Test with whitespace after complete name
+        let next_tokens = token.append(' ');
+        assert!(!next_tokens.is_empty(), "Should accept whitespace after complete name");
+        
+        match &next_tokens[0] {
+            Token::ElementCloseName(_) => (),
+            _ => panic!("Expected ElementCloseName, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_close_name_invalid() {
+        let element = create_test_element();
+        let token = ElementCloseName::new(Arc::clone(&element), "r".to_string());
+        
+        // Test with invalid next character
+        let next_tokens = token.append('x');
+        assert!(next_tokens.is_empty(), "Should reject invalid character");
+    }
 }

lib/xml_schema_validator/src/token/element_close_start.rs

@@ -34,4 +34,76 @@ impl ElementCloseStart {
             }
         }
     }
+}
+
+#[cfg(test)]
+mod element_close_start_tests {
+    use super::*;
+    use std::sync::Arc;
+    use crate::schema;
+
+    fn create_test_element() -> Arc<schema::XsElement> {
+        let element = schema::XsElement {
+            name: "reasoning".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: Some("true".to_string()),
+                text: None,
+                xs_attribute: None,
+                xs_sequence: Some(schema::XsSequence {
+                    text: None,
+                    xs_any: schema::XsAny {
+                        min_occurs: "0".to_string(),
+                        max_occurs: "unbounded".to_string(),
+                        process_contents: "skip".to_string(),
+                    },
+                }),
+            },
+        };
+        Arc::new(element)
+    }
+
+    #[test]
+    fn test_close_start_slash() {
+        let element = create_test_element();
+        let token = ElementCloseStart::new(Arc::clone(&element));
+        
+        // Test with slash
+        let next_tokens = token.append('/');
+        assert!(!next_tokens.is_empty(), "Should accept '/'");
+        
+        match &next_tokens[0] {
+            Token::ElementCloseStart(start) => {
+                assert!(start.has_slash, "Should mark as having slash");
+            },
+            _ => panic!("Expected ElementCloseStart, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_close_start_element_name() {
+        let element = create_test_element();
+        let mut token = ElementCloseStart::new(Arc::clone(&element));
+        token.has_slash = true;
+        
+        // Test with first character of element name
+        let next_tokens = token.append('r');
+        assert!(!next_tokens.is_empty(), "Should accept first letter of element name");
+        
+        match &next_tokens[0] {
+            Token::ElementCloseName(_) => (),
+            _ => panic!("Expected ElementCloseName, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_close_start_invalid_name() {
+        let element = create_test_element();
+        let mut token = ElementCloseStart::new(Arc::clone(&element));
+        token.has_slash = true;
+        
+        // Test with invalid first character of element name
+        let next_tokens = token.append('x');
+        assert!(next_tokens.is_empty(), "Should reject invalid first letter");
+    }
 }

lib/xml_schema_validator/src/token/element_open_end.rs

@@ -15,7 +15,7 @@ impl ElementOpenEnd {
     }
 
     pub fn append(self, c: char) -> Vec<Token> {
-        // Check if the element allows text content (mixed="true")
+        // Check if the element is a mixed content element that allows text
         let allows_text = self.element.xs_complex_type.mixed
             .as_ref()
             .map(|mixed| mixed == "true")
@@ -23,14 +23,140 @@ impl ElementOpenEnd {
 
         if allows_text {
             if c.is_whitespace() {
-                vec![Token::Whitespace(Whitespace::new(
-                    vec![Token::TextContent(TextContent::new(Arc::clone(&self.element)))]
-                ))]
+                // Handle whitespace in mixed content
+                vec![
+                    Token::TextContent(TextContent::new(Arc::clone(&self.element))),
+                    Token::ElementCloseStart(ElementCloseStart::new(Arc::clone(&self.element)))
+                ]
+            } else if c == '<' {
+                // Start of a closing tag or nested element
+                vec![Token::ElementCloseStart(ElementCloseStart::new(Arc::clone(&self.element)))]
             } else {
+                // Text content
                 vec![Token::TextContent(TextContent::new(Arc::clone(&self.element)))]
             }
         } else {
-            vec![]
+            // Element doesn't allow text content, only expect closing tag
+            if c == '<' {
+                vec![Token::ElementCloseStart(ElementCloseStart::new(Arc::clone(&self.element)))]
+            } else if c.is_whitespace() {
+                vec![Token::ElementCloseStart(ElementCloseStart::new(Arc::clone(&self.element)))]
+            } else {
+                // Unexpected content in non-mixed element
+                vec![]
+            }
         }
     }
+}
+
+#[cfg(test)]
+mod element_open_end_tests {
+    use super::*;
+    use std::sync::Arc;
+    use crate::schema;
+
+    fn create_test_element() -> Arc<schema::XsElement> {
+        let element = schema::XsElement {
+            name: "reasoning".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: Some("true".to_string()),
+                text: None,
+                xs_attribute: None,
+                xs_sequence: Some(schema::XsSequence {
+                    text: None,
+                    xs_any: schema::XsAny {
+                        min_occurs: "0".to_string(),
+                        max_occurs: "unbounded".to_string(),
+                        process_contents: "skip".to_string(),
+                    },
+                }),
+            },
+        };
+        Arc::new(element)
+    }
+
+    fn create_element_non_mixed() -> Arc<schema::XsElement> {
+        let element = schema::XsElement {
+            name: "element".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: None, // Not mixed content
+                text: None,
+                xs_attribute: None,
+                xs_sequence: Some(schema::XsSequence {
+                    text: None,
+                    xs_any: schema::XsAny {
+                        min_occurs: "0".to_string(),
+                        max_occurs: "unbounded".to_string(),
+                        process_contents: "skip".to_string(),
+                    },
+                }),
+            },
+        };
+        Arc::new(element)
+    }
+
+    #[test]
+    fn test_append_text_in_mixed_content() {
+        let element = create_test_element();
+        let token = ElementOpenEnd::new(Arc::clone(&element));
+        
+        // Test appending text in mixed content element
+        let next_tokens = token.append('a');
+        assert!(!next_tokens.is_empty(), "Should allow text content in mixed element");
+        
+        match &next_tokens[0] {
+            Token::TextContent(_) => (),
+            _ => panic!("Expected TextContent, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_append_closing_tag_start() {
+        let element = create_test_element();
+        let token = ElementOpenEnd::new(Arc::clone(&element));
+        
+        // Test starting a closing tag
+        let next_tokens = token.append('<');
+        assert!(!next_tokens.is_empty(), "Should allow closing tag start");
+        
+        match &next_tokens[0] {
+            Token::ElementCloseStart(_) => (),
+            _ => panic!("Expected ElementCloseStart, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_append_whitespace_in_mixed() {
+        let element = create_test_element();
+        let token = ElementOpenEnd::new(Arc::clone(&element));
+        
+        // Test whitespace in mixed content
+        let next_tokens = token.append(' ');
+        assert!(!next_tokens.is_empty(), "Should allow whitespace in mixed content");
+        
+        match &next_tokens[0] {
+            Token::TextContent(_) => (),
+            _ => panic!("Expected TextContent, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_append_in_non_mixed() {
+        let element = create_element_non_mixed();
+        let token = ElementOpenEnd::new(Arc::clone(&element));
+        
+        // Test appending text in non-mixed content (should reject)
+        let next_tokens = token.clone().append('a');
+        assert!(next_tokens.is_empty(), "Should reject text in non-mixed element");
+        
+        // Test whitespace (should accept leading to closing tag)
+        let next_tokens = token.clone().append(' ');
+        assert!(!next_tokens.is_empty(), "Should allow whitespace in non-mixed");
+        
+        // Test closing tag start
+        let next_tokens = token.append('<');
+        assert!(!next_tokens.is_empty(), "Should allow closing tag start in non-mixed");
+    }
 }

lib/xml_schema_validator/src/token/element_open_name.rs

@@ -16,25 +16,218 @@ impl ElementOpenName {
                 element,
             })
         } else {
-            Err(Error::NotImplemented)
+            Err(Error::InvalidXml(format!("Element name '{}' doesn't match '{}'", buffer, element.name)))
         }
     }
 
     pub fn append(self, c: char) -> Vec<Token> {
-        let buffer = self.buffer + &c.to_string();
-        if self.element.name.starts_with(buffer.as_str()) {
-            vec![Token::ElementName(ElementOpenName {
-                buffer,
-                element: Arc::clone(&self.element),
-            })]
-        } else if c.is_whitespace() {
-            vec![Token::Whitespace(Whitespace::new(
-                vec![Token::ElementOpen(ElementOpenEnd::new(Arc::clone(&self.element)))]
-            ))]
-        } else if '>' == c {
-            vec![Token::ElementOpen(ElementOpenEnd::new(Arc::clone(&self.element)))]
+        if c.is_whitespace() {
+            // Check if we've matched the full element name
+            if self.element.name == self.buffer {
+                // Element name is complete, handle whitespace after name
+                
+                // Check if the element has attributes
+                if let Some(attributes) = &self.element.xs_complex_type.xs_attribute {
+                    if !attributes.is_empty() {
+                        // Transition to attribute name parsing
+                        return vec![Token::AttributeName(AttributeName::new(Arc::clone(&self.element)))];
+                    }
+                }
+                
+                // No attributes, continue to element open or self-close
+                return vec![
+                    Token::ElementOpen(ElementOpenEnd::new(Arc::clone(&self.element))),
+                    Token::ElementSelfClose(ElementSelfClose::new())
+                ];
+            } else {
+                // Whitespace in the middle of an element name is invalid
+                return vec![];
+            }
+        } else if c == '>' {
+            // Check if we've matched the full element name
+            if self.element.name == self.buffer {
+                // Element name is complete, end of opening tag
+                return vec![Token::ElementOpen(ElementOpenEnd::new(Arc::clone(&self.element)))];
+            } else {
+                // '>' before full element name is invalid
+                return vec![];
+            }
+        } else if c == '/' {
+            // Check if we've matched the full element name
+            if self.element.name == self.buffer {
+                // Element name is complete, self-closing tag
+                return vec![Token::ElementSelfClose(ElementSelfClose::new())];
+            } else {
+                // '/' before full element name is invalid
+                return vec![];
+            }
         } else {
-            vec![]
+            // Add character to buffer and check if we're still building a valid element name
+            let new_buffer = self.buffer.clone() + &c.to_string();
+            
+            if self.element.name.starts_with(&new_buffer) {
+                // Still building the element name
+                return vec![Token::ElementName(ElementOpenName {
+                    buffer: new_buffer,
+                    element: Arc::clone(&self.element),
+                })];
+            } else {
+                // Character doesn't match what we expect for this element name
+                return vec![];
+            }
         }
     }
+}
+
+#[cfg(test)]
+mod element_open_name_tests {
+    use super::*;
+    use std::sync::Arc;
+    use crate::schema;
+
+    fn create_test_element() -> Arc<schema::XsElement> {
+        let element = schema::XsElement {
+            name: "reasoning".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: Some("true".to_string()),
+                text: None,
+                xs_attribute: None,
+                xs_sequence: Some(schema::XsSequence {
+                    text: None,
+                    xs_any: schema::XsAny {
+                        min_occurs: "0".to_string(),
+                        max_occurs: "unbounded".to_string(),
+                        process_contents: "skip".to_string(),
+                    },
+                }),
+            },
+        };
+        Arc::new(element)
+    }
+
+    fn create_element_with_attributes() -> Arc<schema::XsElement> {
+        let element = schema::XsElement {
+            name: "delete".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: None,
+                text: None,
+                xs_attribute: Some(vec![
+                    schema::XsAttribute {
+                        name: "id".to_string(),
+                        xs_attribute_type: "xs:string".to_string(),
+                        xs_attribute_use: "required".to_string(),
+                    }
+                ]),
+                xs_sequence: None,
+            },
+        };
+        Arc::new(element)
+    }
+
+    #[test]
+    fn test_append_character_by_character() {
+        let element = create_test_element();
+        
+        // Test each character in "reasoning"
+        let mut buffer = String::new();
+        for c in "reasoning".chars() {
+            buffer.push(c);
+            let result = ElementOpenName::new(Arc::clone(&element), buffer.clone());
+            assert!(result.is_ok(), "Failed to create ElementOpenName with buffer '{}'", buffer);
+        }
+    }
+
+    #[test]
+    fn test_append_full_name() {
+        let element = create_test_element();
+        
+        // Test with the full element name
+        let result = ElementOpenName::new(Arc::clone(&element), "reasoning".to_string());
+        assert!(result.is_ok(), "Failed to create ElementOpenName with full element name");
+        
+        let element_name = result.unwrap();
+        
+        // Test transition after full name with '>'
+        let next_tokens = element_name.append('>');
+        assert!(!next_tokens.is_empty(), "Should accept '>' after full element name");
+        
+        match &next_tokens[0] {
+            Token::ElementOpen(_) => (),
+            _ => panic!("Expected ElementOpen, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_append_space_after_name() {
+        let element = create_test_element();
+        
+        // Test with the full element name
+        let element_name = ElementOpenName::new(Arc::clone(&element), "reasoning".to_string()).unwrap();
+        
+        // Test transition after full name with space
+        let next_tokens = element_name.append(' ');
+        assert!(!next_tokens.is_empty(), "Should accept space after full element name");
+        
+        match &next_tokens[0] {
+            Token::ElementOpen(_) => (),
+            _ => panic!("Expected ElementOpen, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_append_with_attributes() {
+        let element = create_element_with_attributes();
+        
+        // Test with the full element name
+        let element_name = ElementOpenName::new(Arc::clone(&element), "delete".to_string()).unwrap();
+        
+        // Test transition after full name with space (expecting attribute name)
+        let next_tokens = element_name.append(' ');
+        assert!(!next_tokens.is_empty(), "Should accept space after element name with attributes");
+        
+        // The space should lead to a whitespace token, which should then lead to attribute name
+        match &next_tokens[0] {
+            Token::AttributeName(_) => (),
+            _ => panic!("Expected AttributeName, got {:?}", next_tokens[0]),
+        }
+    }
+    
+    #[test]
+    fn test_append_slash_after_name() {
+        let element = create_test_element();
+        
+        // Test with the full element name
+        let element_name = ElementOpenName::new(Arc::clone(&element), "reasoning".to_string()).unwrap();
+        
+        // Test transition after full name with /
+        let next_tokens = element_name.append('/');
+        assert!(!next_tokens.is_empty(), "Should accept '/' after full element name");
+        
+        match &next_tokens[0] {
+            Token::ElementSelfClose(_) => (),
+            _ => panic!("Expected ElementSelfClose, got {:?}", next_tokens[0]),
+        }
+    }
+    
+    #[test]
+    fn test_invalid_continuation() {
+        let element = create_test_element();
+        
+        // Test with a partial element name and an invalid continuation character
+        let element_name = ElementOpenName::new(Arc::clone(&element), "reas".to_string()).unwrap();
+        let next_tokens = element_name.append('x');
+        assert!(next_tokens.is_empty(), "Should reject invalid character in element name");
+        
+        // Test whitespace in middle of name - need to create a new instance since append consumes self
+        let element_name = ElementOpenName::new(Arc::clone(&element), "reas".to_string()).unwrap();
+        let next_tokens = element_name.append(' ');
+        assert!(next_tokens.is_empty(), "Should reject whitespace in middle of element name");
+        
+        // Test > in middle of name - need to create a new instance since append consumes self
+        let element_name = ElementOpenName::new(Arc::clone(&element), "reas".to_string()).unwrap();
+        let next_tokens = element_name.append('>');
+        assert!(next_tokens.is_empty(), "Should reject '>' in middle of element name");
+    }
 }

lib/xml_schema_validator/src/token/element_self_close.rs

@@ -0,0 +1,68 @@
+use crate::*;
+
+/// Represents a self-closing XML element
+#[derive(Clone, Debug)]
+pub struct ElementSelfClose {
+}
+
+impl ElementSelfClose {
+    /// Create a new self-closing element
+    pub fn new() -> Self {
+        Self {
+        }
+    }
+
+    /// Append a character to the self-closing tag and return possible continuations
+    pub fn append(self, c: char) -> Vec<Token> {
+        if c == '>' {
+            vec![Token::EndOfFile(EndOfFile::new())]
+        } else if c.is_whitespace() {
+            vec![Token::ElementSelfClose(self)]
+        } else {
+            vec![]
+        }
+    }
+}
+
+#[cfg(test)]
+mod element_self_close_tests {
+    use super::*;
+
+    #[test]
+    fn test_self_close_with_greater_than() {
+        let token = ElementSelfClose::new();
+        
+        // Test closing with >
+        let next_tokens = token.append('>');
+        assert!(!next_tokens.is_empty(), "Should accept '>' to close self-closing tag");
+        assert_eq!(next_tokens.len(), 1, "Should have exactly one token after '>'");
+        
+        match &next_tokens[0] {
+            Token::EndOfFile(_) => (),
+            _ => panic!("Expected EndOfFile, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_self_close_with_whitespace() {
+        let token = ElementSelfClose::new();
+        
+        // Test with whitespace
+        let next_tokens = token.append(' ');
+        assert!(!next_tokens.is_empty(), "Should accept whitespace in self-closing tag");
+        
+        match &next_tokens[0] {
+            Token::ElementSelfClose(_) => (),
+            _ => panic!("Expected ElementSelfClose, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_self_close_with_invalid_char() {
+        let token = ElementSelfClose::new();
+        
+        // Test with invalid character
+        let next_tokens = token.append('a');
+        assert!(next_tokens.is_empty(), "Should reject invalid character in self-closing tag");
+    }
+}

lib/xml_schema_validator/src/token/end_of_file.rs

@@ -13,4 +13,18 @@ impl EndOfFile {
     pub fn append(self, _c: char) -> Vec<Token> {
         vec![]
     }
+}
+
+#[cfg(test)]
+mod end_of_file_tests {
+    use super::*;
+
+    #[test]
+    fn test_end_of_file_any_char() {
+        let token = EndOfFile::new();
+        
+        // Test with any character (should reject)
+        let next_tokens = token.append('a');
+        assert!(next_tokens.is_empty(), "Should reject any character at EOF");
+    }
 }

lib/xml_schema_validator/src/token/mod.rs

@@ -1,4 +1,3 @@
-mod element_close_end;
 mod element_close_name;
 mod element_close_start;
 mod element_open_name;
@@ -7,9 +6,10 @@ mod element_open_start;
 mod end_of_file;
 mod start_of_file;
 mod text_content;
-mod whitespace;
+mod attribute_name;
+mod attribute_value;
+mod element_self_close;
 
-pub use element_close_end::ElementCloseEnd;
 pub use element_close_name::ElementCloseName;
 pub use element_close_start::ElementCloseStart;
 pub use element_open_name::ElementOpenName;
@@ -18,12 +18,13 @@ pub use element_open_start::ElementOpenStart;
 pub use end_of_file::EndOfFile;
 pub use start_of_file::StartOfFile;
 pub use text_content::TextContent;
-pub use whitespace::Whitespace;
+pub use attribute_name::AttributeName;
+pub use attribute_value::AttributeValue;
+pub use element_self_close::ElementSelfClose;
 
 /// Represents the different states of XML token parsing
 #[derive(Clone, Debug)]
 pub enum Token {
-    ElementCloseEnd(ElementCloseEnd),
     ElementCloseName(ElementCloseName),
     ElementCloseStart(ElementCloseStart),
     ElementName(ElementOpenName),
@@ -32,13 +33,14 @@ pub enum Token {
     EndOfFile(EndOfFile),
     StartOfFile(StartOfFile),
     TextContent(TextContent),
-    Whitespace(Whitespace),
+    AttributeName(AttributeName),
+    AttributeValue(AttributeValue),
+    ElementSelfClose(ElementSelfClose),
 }
 
 impl Token {
     pub fn append(self, c: char) -> Vec<Token> {
         match self {
-            Token::ElementCloseEnd(element_close_end) => element_close_end.append(c),
             Token::ElementCloseName(element_close_name) => element_close_name.append(c),
             Token::ElementCloseStart(element_close_start) => element_close_start.append(c),
             Token::ElementName(element_name) => element_name.append(c),
@@ -47,7 +49,9 @@ impl Token {
             Token::EndOfFile(end_of_file) => end_of_file.append(c),
             Token::StartOfFile(start_of_file) => start_of_file.append(c),
             Token::TextContent(text_content) => text_content.append(c),
-            Token::Whitespace(whitespace) => whitespace.append(c),
+            Token::AttributeName(attribute_name) => attribute_name.append(c),
+            Token::AttributeValue(attribute_value) => attribute_value.append(c),
+            Token::ElementSelfClose(element_self_close) => element_self_close.append(c),
         }
     }
 
@@ -57,4 +61,232 @@ impl Token {
             _ => false,
         }
     }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::sync::Arc;
+    use crate::schema;
+
+    fn create_test_element() -> Arc<schema::XsElement> {
+        let element = schema::XsElement {
+            name: "reasoning".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: Some("true".to_string()),
+                text: None,
+                xs_attribute: None,
+                xs_sequence: Some(schema::XsSequence {
+                    text: None,
+                    xs_any: schema::XsAny {
+                        min_occurs: "0".to_string(),
+                        max_occurs: "unbounded".to_string(),
+                        process_contents: "skip".to_string(),
+                    },
+                }),
+            },
+        };
+        Arc::new(element)
+    }
+
+    #[test]
+    fn test_start_of_file_transition() {
+        let element = create_test_element();
+        let token = Token::StartOfFile(StartOfFile::new(Arc::clone(&element)));
+        
+        // Should only accept '<'
+        let next_tokens = token.append('<');
+        assert_eq!(next_tokens.len(), 1, "Should have one transition");
+        
+        match &next_tokens[0] {
+            Token::ElementStart(_) => (),
+            _ => panic!("Expected ElementStart, got {:?}", next_tokens[0]),
+        }
+        
+        // Should reject other characters
+        let token = Token::StartOfFile(StartOfFile::new(Arc::clone(&element)));
+        let next_tokens = token.append('a');
+        assert_eq!(next_tokens.len(), 0, "Should reject 'a'");
+    }
+    
+    #[test]
+    fn test_element_start_transition() {
+        let element = create_test_element();
+        let token = Token::ElementStart(ElementOpenStart::new(Arc::clone(&element)));
+        
+        // Should accept first char of element name
+        let next_tokens = token.append('r');
+        assert_eq!(next_tokens.len(), 1, "Should have one transition");
+        
+        match &next_tokens[0] {
+            Token::ElementName(_) => (),
+            _ => panic!("Expected ElementName, got {:?}", next_tokens[0]),
+        }
+        
+        // Should reject invalid char
+        let token = Token::ElementStart(ElementOpenStart::new(Arc::clone(&element)));
+        let next_tokens = token.append('x');
+        assert_eq!(next_tokens.len(), 0, "Should reject 'x'");
+    }
+    
+    #[test]
+    fn test_element_name_transition() {
+        let element = create_test_element();
+        let token = Token::ElementName(ElementOpenName::new(Arc::clone(&element), "r".to_string()).unwrap());
+        
+        // Continue building name
+        let next_tokens = token.append('e');
+        assert_eq!(next_tokens.len(), 1, "Should have one transition");
+        
+        match &next_tokens[0] {
+            Token::ElementName(_) => (),
+            _ => panic!("Expected ElementName, got {:?}", next_tokens[0]),
+        }
+        
+        // Build complete name and end with >
+        // Remove the unused buffer
+        let mut current_tokens = vec![Token::ElementStart(ElementOpenStart::new(Arc::clone(&element)))];
+        
+        for c in "reasoning>".chars() {
+            let mut new_tokens = Vec::new();
+            for token in current_tokens {
+                new_tokens.append(&mut token.append(c));
+            }
+            current_tokens = new_tokens;
+            
+            assert!(!current_tokens.is_empty(), 
+                "Should have valid transition after appending '{}'", c);
+        }
+        
+        // Check final state
+        assert_eq!(current_tokens.len(), 1, "Should have one final transition");
+        match &current_tokens[0] {
+            Token::ElementOpen(_) => (),
+            _ => panic!("Expected ElementOpen, got {:?}", current_tokens[0]),
+        }
+    }
+    
+    #[test]
+    fn test_element_open_end_transition() {
+        let element = create_test_element();
+        let token = Token::ElementOpen(ElementOpenEnd::new(Arc::clone(&element)));
+        
+        // Should accept text content in mixed element
+        let next_tokens = token.append('t');
+        assert_eq!(next_tokens.len(), 1, "Should accept text");
+        
+        match &next_tokens[0] {
+            Token::TextContent(_) => (),
+            _ => panic!("Expected TextContent, got {:?}", next_tokens[0]),
+        }
+        
+        // Should accept closing tag start
+        let token = Token::ElementOpen(ElementOpenEnd::new(Arc::clone(&element)));
+        let next_tokens = token.append('<');
+        assert_eq!(next_tokens.len(), 1, "Should accept '<'");
+        
+        match &next_tokens[0] {
+            Token::ElementCloseStart(_) => (),
+            _ => panic!("Expected ElementCloseStart, got {:?}", next_tokens[0]),
+        }
+    }
+    
+    #[test]
+    fn test_text_content_transition() {
+        let element = create_test_element();
+        let token = Token::TextContent(TextContent::new(Arc::clone(&element)));
+        
+        // Should accept more text
+        let next_tokens = token.append('a');
+        assert_eq!(next_tokens.len(), 1, "Should accept more text");
+        
+        match &next_tokens[0] {
+            Token::TextContent(_) => (),
+            _ => panic!("Expected TextContent, got {:?}", next_tokens[0]),
+        }
+        
+        // Should accept closing tag start
+        let token = Token::TextContent(TextContent::new(Arc::clone(&element)));
+        let next_tokens = token.append('<');
+        assert_eq!(next_tokens.len(), 1, "Should accept '<'");
+        
+        match &next_tokens[0] {
+            Token::ElementCloseStart(_) => (),
+            _ => panic!("Expected ElementCloseStart, got {:?}", next_tokens[0]),
+        }
+    }
+    
+    #[test]
+    fn test_element_close_start_transition() {
+        let element = create_test_element();
+        let token = Token::ElementCloseStart(ElementCloseStart::new(Arc::clone(&element)));
+        
+        // Should accept /
+        let next_tokens = token.append('/');
+        assert_eq!(next_tokens.len(), 1, "Should accept '/'");
+        
+        let token = next_tokens[0].clone();
+        match token {
+            Token::ElementCloseStart(_) => (),
+            _ => panic!("Expected ElementCloseStart with slash, got {:?}", token),
+        }
+        
+        // After /, should accept first char of element name
+        let next_tokens = token.append('r');
+        assert_eq!(next_tokens.len(), 1, "Should accept 'r'");
+        
+        match &next_tokens[0] {
+            Token::ElementCloseName(_) => (),
+            _ => panic!("Expected ElementCloseName, got {:?}", next_tokens[0]),
+        }
+    }
+    
+    #[test]
+    fn test_element_close_name_transition() {
+        let element = create_test_element();
+        let token = Token::ElementCloseStart(ElementCloseStart::new(Arc::clone(&element)));
+        
+        // Add / first
+        let next_tokens = token.append('/');
+        assert_eq!(next_tokens.len(), 1, "Should accept '/'");
+        
+        // Start building name - clone the token to avoid moving out of the vector
+        let next_tokens = next_tokens[0].clone().append('r');
+        assert_eq!(next_tokens.len(), 1, "Should accept 'r'");
+        
+        // Continue with each character
+        let mut current_tokens = next_tokens;
+        for c in "easoning".chars() {
+            let mut new_tokens = Vec::new();
+            for token in current_tokens {
+                new_tokens.append(&mut token.append(c));
+            }
+            current_tokens = new_tokens;
+            
+            assert!(!current_tokens.is_empty(), 
+                "Should have valid transition after appending '{}'", c);
+        }
+        
+        // Final > to close - clone the token to avoid moving out of the vector
+        let next_tokens = current_tokens[0].clone().append('>');
+        assert_eq!(next_tokens.len(), 1, "Should accept '>'");
+        
+        match &next_tokens[0] {
+            Token::EndOfFile(_) => (),
+            _ => panic!("Expected EndOfFile, got {:?}", next_tokens[0]),
+        }
+    }
+    
+    #[test]
+    fn test_end_of_file_properties() {
+        let token = Token::EndOfFile(EndOfFile::new());
+        
+        // Should be identified as EOF
+        assert!(token.is_eof(), "EndOfFile should be recognized as EOF");
+        
+        // Should not accept any more input
+        let next_tokens = token.append('a');
+        assert_eq!(next_tokens.len(), 0, "Should not accept any more input");
+    }
 }

lib/xml_schema_validator/src/token/start_of_file.rs

@@ -19,4 +19,57 @@ impl StartOfFile {
             vec![]
         }
     }
+}
+
+#[cfg(test)]
+mod start_of_file_tests {
+    use super::*;
+    use std::sync::Arc;
+    use crate::schema;
+
+    fn create_test_element() -> Arc<schema::XsElement> {
+        let element = schema::XsElement {
+            name: "reasoning".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: Some("true".to_string()),
+                text: None,
+                xs_attribute: None,
+                xs_sequence: Some(schema::XsSequence {
+                    text: None,
+                    xs_any: schema::XsAny {
+                        min_occurs: "0".to_string(),
+                        max_occurs: "unbounded".to_string(),
+                        process_contents: "skip".to_string(),
+                    },
+                }),
+            },
+        };
+        Arc::new(element)
+    }
+
+    #[test]
+    fn test_start_of_file_less_than() {
+        let element = create_test_element();
+        let token = StartOfFile::new(Arc::clone(&element));
+        
+        // Test with <
+        let next_tokens = token.append('<');
+        assert!(!next_tokens.is_empty(), "Should accept '<'");
+        
+        match &next_tokens[0] {
+            Token::ElementStart(_) => (),
+            _ => panic!("Expected ElementStart, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_start_of_file_invalid() {
+        let element = create_test_element();
+        let token = StartOfFile::new(Arc::clone(&element));
+        
+        // Test with invalid character
+        let next_tokens = token.append('a');
+        assert!(next_tokens.is_empty(), "Should reject invalid character");
+    }
 }

lib/xml_schema_validator/src/token/text_content.rs

@@ -23,4 +23,62 @@ impl TextContent {
             ))]
         }
     }
+}
+
+#[cfg(test)]
+mod text_content_tests {
+    use super::*;
+    use std::sync::Arc;
+    use crate::schema;
+
+    fn create_test_element() -> Arc<schema::XsElement> {
+        let element = schema::XsElement {
+            name: "reasoning".to_string(),
+            text: None,
+            xs_complex_type: schema::XsComplexType {
+                mixed: Some("true".to_string()),
+                text: None,
+                xs_attribute: None,
+                xs_sequence: Some(schema::XsSequence {
+                    text: None,
+                    xs_any: schema::XsAny {
+                        min_occurs: "0".to_string(),
+                        max_occurs: "unbounded".to_string(),
+                        process_contents: "skip".to_string(),
+                    },
+                }),
+            },
+        };
+        Arc::new(element)
+    }
+
+    #[test]
+    fn test_text_content_normal_char() {
+        let element = create_test_element();
+        let token = TextContent::new(Arc::clone(&element));
+        
+        // Test with normal character
+        let next_tokens = token.append('a');
+        assert!(!next_tokens.is_empty(), "Should accept normal character");
+        
+        match &next_tokens[0] {
+            Token::TextContent(_) => (),
+            _ => panic!("Expected TextContent, got {:?}", next_tokens[0]),
+        }
+    }
+
+    #[test]
+    fn test_text_content_less_than() {
+        let element = create_test_element();
+        let token = TextContent::new(Arc::clone(&element));
+        
+        // Test with < (start of tag)
+        let next_tokens = token.append('<');
+        assert!(!next_tokens.is_empty(), "Should accept '<'");
+        
+        match &next_tokens[0] {
+            Token::ElementCloseStart(_) => (),
+            _ => panic!("Expected ElementCloseStart, got {:?}", next_tokens[0]),
+        }
+    }
 }

lib/xml_schema_validator/src/token/whitespace.rs

@@ -1,29 +0,0 @@
-use crate::*;
-
-/// Represents one or more whitespace characters
-#[derive(Clone, Debug)]
-pub struct Whitespace {
-    next: Vec<Token>
-}
-
-impl Whitespace {
-    pub fn new(next: Vec<Token>) -> Self {
-        Self {
-            next
-        }
-    }
-
-    pub fn append(self, c: char) -> Vec<Token> {
-        if c.is_whitespace() {
-            vec![Token::Whitespace(Self {
-                next: self.next
-            })]
-        } else {
-            let mut tokens = vec![];
-            for token in self.next {
-                tokens.append(&mut token.append(c))
-            }
-            tokens
-        }
-    }
-}

lib/xml_schema_validator/tests/test_logits_processor.py

@@ -1,14 +1,14 @@
-from xml_logits_processor import XmlLogitsProcessor
-from threading import Thread
-from transformers import pipeline, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoTokenizer
+from xml_schema_validator import XmlLogitsProcessor
 import os
-import sys
+import time
+import torch
 
-def test_logits_processor():
-    print("test_logits_processor")
+model_path = "../../model"
+api_token = os.environ["SIA_HF_API_KEY"]
 
-    model_name = "../../model"
-    xml_schema_text = """<?xml version="1.0" encoding="UTF-8"?>
+xml_schema_actions = open("../../action_schema.xsd").read()
+xml_schema_only_root_node = """<?xml version="1.0" encoding="UTF-8"?>
 <xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema" elementFormDefault="qualified">
   <xs:element name="root">
     <xs:complexType mixed="true">
@@ -19,53 +19,106 @@ def test_logits_processor():
   </xs:element>
 </xs:schema>"""
 
-    model = "google/gemma-3-1b-it"
+def test_logits_processor_init_copy():
+    """Test the initialization of the LogitsProcessor"""
 
-    pipline = pipeline(
-        "text-generation",
-        model=model,
-        token=os.environ["SIA_HF_API_KEY"],
-    )
+    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+    logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node)
+    assert logits_processor.core is not None
 
-    tokenizer = AutoTokenizer.from_pretrained(
-        model,
-        token=os.environ["SIA_HF_API_KEY"],
-    )
+    logits_processor_copy = logits_processor.copy()
+    assert logits_processor_copy.core is not None
 
-    logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_text)
-                        
-    messages = [
-        {"role": "system", "content": "Always respond with <root>message by the user</root>. So if the user says 'hello world', you response <root>hello world</root>"},
-        {"role": "user", "content": "hello, I am the user"}
-    ]
+def test_xml_schema_parsing():
+    """Test basic XML schema parsing with different element types"""
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+    logits_processor = XmlLogitsProcessor(tokenizer, xml_schema_only_root_node)
+    assert logits_processor.core is not None
 
-    text_inputs = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
+    element_names = logits_processor.core.get_element_names()
+    assert len(element_names) > 0
+    assert "root" in element_names
 
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=True
-    )
+def test_basic_xml_validation():
+    """Test basic validation of XML fragments against the schema"""
 
-    generation_kwargs = {
-        "text_inputs": text_inputs,
-        "max_new_tokens": 20,
-        "streamer": streamer,
-        #"logits_processor": [logits_processor],
-    }
+    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+    processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
+    
+    # Test appending valid XML fragments
+    assert processor.core.copy().append("<reasoning>")
+    assert processor.core.copy().append("<reasoning>Test content</reasoning>")
+    
+    # Test appending invalid XML fragments
+    assert not processor.core.copy().append("<invalid>")
+    assert not processor.core.copy().append("<reasoning></invalid>")
 
-    generation_thread = Thread(
-        target=pipline,
-        kwargs=generation_kwargs
-    )
+def test_token_masking():
+    """Test that invalid tokens are properly masked"""
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+    processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
+    
+    # Create dummy input_ids and scores
+    input_ids = torch.tensor([tokenizer.encode("<reasoning>")])
+    scores = torch.ones((1, len(tokenizer)))  # All tokens have equal probability
+    
+    # Process the scores
+    processed_scores = processor(input_ids, scores)
+    
+    # Verify that valid continuations still have positive scores
+    assert processed_scores[0, tokenizer.encode(" ", add_special_tokens=False)[0]] > -float('inf')
+    
+    # Verify that invalid continuations are masked
+    assert processed_scores[0, tokenizer.encode("<invalid>", add_special_tokens=False)[0]] == -float('inf')
 
-    generation_thread.start()
-
-    for text in streamer:
-        print(text, end="", file=sys.stderr, flush=True)
-
-    generation_thread.join()
+def test_performance():
+    """Test performance with larger XML documents"""
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+    processor = XmlLogitsProcessor(tokenizer, xml_schema_actions)
+    
+    # Generate a larger XML document
+    large_xml = "<reasoning>" + "Test content. " * 100 + "</reasoning>"
+    
+    # Measure time to process
+    start_time = time.time()
+    processor.core.append(large_xml)
+    processing_time = time.time() - start_time
+    
+    print(f"Processing time for large XML: {processing_time:.4f} seconds")
+    # You might want to assert that processing time is below a threshold
 
+def test_subword_tokens():
+    """Test handling of subword tokens that might split XML tags"""
+    tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
+    schema_text = open("../../action_schema.xsd").read()
+    
+    processor = XmlLogitsProcessor(tokenizer, schema_text)
+    
+    # Find some XML tag that gets split into multiple tokens by your tokenizer
+    tag = "<reasoning>"
+    tokens = tokenizer.encode(tag, add_special_tokens=False)
+    
+    if len(tokens) > 1:
+        print(f"Tag '{tag}' is split into {len(tokens)} tokens")
+        
+        # Test that the processor can handle these split tokens
+        import torch
+        input_ids = torch.tensor([[tokens[0]]])
+        scores = torch.ones((1, len(tokenizer)))
+        
+        processed_scores = processor(input_ids, scores)
+        
+        # The next token in the sequence should have a high score
+        assert processed_scores[0, tokens[1]] > -float('inf')
+    
 if __name__ == "__main__":
-    test_logits_processor()
+    # Run individual tests for debugging
+    test_logits_processor_init_copy()
+    test_xml_schema_parsing()
+    test_basic_xml_validation()
+    test_token_masking()
+    test_performance()
+    test_subword_tokens()