poimen-memory/crates/mem-core/src/gate_parser.rs

use thiserror::Error;

/// Parsed gate response from the model.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GateResponse {
    pub think: String,
    pub update_gate: bool,
    pub candidate: String,
    pub exit_gate: bool,
}

/// Parse error with context.
#[derive(Error, Debug, Clone)]
#[error("Parse error in {tag}: {message}\nRaw: {raw}")]
pub struct ParseError {
    pub tag: String,
    pub message: String,
    pub raw: String,
}

/// Parse a gate response from model output.
pub fn parse_gate_response(response: &str) -> Result<GateResponse, ParseError> {
    // Extract <think>...</think> — last one before first <check>
    let think = extract_think(response)?;
    
    // Extract <check>yes|no</check>
    let check_value = extract_tag_value(response, "check")?;
    let update_gate = match check_value.trim().to_lowercase().as_str() {
        "yes" => true,
        "no" => false,
        _ => {
            return Err(ParseError {
                tag: "check".to_string(),
                message: format!("must be 'yes' or 'no', got '{}'", check_value),
                raw: truncate(response, 200),
            });
        }
    };
    
    // Extract <update>...</update>
    let candidate = extract_tag_value(response, "update")?;
    
    // Extract <next>continue|end</next>
    let next_value = extract_tag_value(response, "next")?;
    let exit_gate = match next_value.trim().to_lowercase().as_str() {
        "continue" => false,
        "end" => true,
        _ => {
            return Err(ParseError {
                tag: "next".to_string(),
                message: format!("must be 'continue' or 'end', got '{}'", next_value),
                raw: truncate(response, 200),
            });
        }
    };
    
    Ok(GateResponse {
        think,
        update_gate,
        candidate,
        exit_gate,
    })
}

/// Extract the last <think>...</think> before first <check>.
fn extract_think(response: &str) -> Result<String, ParseError> {
    let check_pos = response.find("<check>").ok_or_else(|| ParseError {
        tag: "check".to_string(),
        message: "tag not found".to_string(),
        raw: truncate(response, 200),
    })?;
    
    // Look for the last </think> before the <check>
    let before_check = &response[..check_pos];
    if let Some(end_pos) = before_check.rfind("</think>") {
        // Look for the last <think> before this </think>
        if let Some(start_pos) = before_check[..end_pos].rfind("<think>") {
            let think_content = &before_check[start_pos + 7..end_pos]; // 7 = "<think>".len()
            return Ok(think_content.to_string());
        }
    }
    
    Err(ParseError {
        tag: "think".to_string(),
        message: "tag not found or not properly closed".to_string(),
        raw: truncate(response, 200),
    })
}

/// Extract content between <tag>...</tag>, ensuring it appears exactly once.
fn extract_tag_value(response: &str, tag: &str) -> Result<String, ParseError> {
    let open_tag = format!("<{}>", tag);
    let close_tag = format!("</{}>", tag);
    
    // Check if tag appears at all
    if !response.contains(&open_tag) {
        return Err(ParseError {
            tag: tag.to_string(),
            message: "tag not found".to_string(),
            raw: truncate(response, 200),
        });
    }
    
    // Check for duplicates
    let open_count = response.matches(&open_tag).count();
    let close_count = response.matches(&close_tag).count();
    
    if open_count > 1 || close_count > 1 {
        return Err(ParseError {
            tag: tag.to_string(),
            message: format!(
                "tag appears {} times (expected exactly 1)",
                open_count.max(close_count)
            ),
            raw: truncate(response, 200),
        });
    }
    
    if close_count == 0 {
        return Err(ParseError {
            tag: tag.to_string(),
            message: "tag not properly closed".to_string(),
            raw: truncate(response, 200),
        });
    }
    
    // Extract content
    let start_idx = response.find(&open_tag).unwrap() + open_tag.len();
    let end_idx = response.find(&close_tag).unwrap();
    
    if start_idx > end_idx {
        return Err(ParseError {
            tag: tag.to_string(),
            message: "malformed tag structure".to_string(),
            raw: truncate(response, 200),
        });
    }
    
    Ok(response[start_idx..end_idx].to_string())
}

/// Truncate a string for display.
fn truncate(s: &str, max_len: usize) -> String {
    if s.len() > max_len {
        format!("{}...", &s[..max_len])
    } else {
        s.to_string()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_wellformed_yes_continue() {
        let response = r#"
<think>This is reasoning</think>
<check>yes</check>
<update>Memory update text</update>
<next>continue</next>
"#;
        let result = parse_gate_response(response).unwrap();
        assert_eq!(result.think, "This is reasoning");
        assert!(result.update_gate);
        assert_eq!(result.candidate, "Memory update text");
        assert!(!result.exit_gate);
    }

    #[test]
    fn test_missing_tag() {
        let response = r#"
<think>This is reasoning</think>
<check>yes</check>
<next>continue</next>
"#;
        let result = parse_gate_response(response);
        assert!(result.is_err());
        assert_eq!(result.unwrap_err().tag, "update");
    }
}