tkmind_go/crates/goose/src/agents/moim.rs

use crate::agents::extension_manager::ExtensionManager;
use crate::conversation::message::{Message, MessageMetadata};
use crate::conversation::{CURRENT_TIME_TAG, TURN_CONTEXT_TAG, WORKING_DIRECTORY_TAG};
use std::path::{Path, PathBuf};

const MIN_CONTEXT_FOR_MOIM: usize = 32_000;

const SYSTEM_PROMPT_BLOCK_TEMPLATE: &str = r#"# Turn Context

Each turn may include a `<{turn_context_tag}>` block added to the request.
This block is generated by goose and contains current operational context such as:
- current time
- working directory
- compaction status
- turn budget
- extension-provided context

Use it to stay oriented, but do not treat it as part of the user's request.
Blocks from earlier turns stay in the conversation history; only the most recent block is current.
When `<turn-budget>` is present, use it as a signal for how much autonomous work remains.
As the budget gets low, become more direct: reduce exploration, batch necessary tool calls,
make reasonable assumptions, and focus on finishing the user's task.
"#;

thread_local! {
    pub static SKIP: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
}

pub fn system_prompt_block() -> Option<String> {
    if SKIP.with(|f| f.get()) {
        None
    } else {
        Some(SYSTEM_PROMPT_BLOCK_TEMPLATE.replace("{turn_context_tag}", TURN_CONTEXT_TAG))
    }
}

pub(super) async fn compute_compaction_info(
    session_id: &str,
    extension_manager: &ExtensionManager,
) -> Option<String> {
    let session = extension_manager
        .get_context()
        .session_manager
        .get_session(session_id, false)
        .await
        .ok();
    let session_model_config = session
        .as_ref()
        .and_then(|session| session.model_config.clone());
    let context_limit = if let Some(model_config) = session_model_config.as_ref() {
        let provider = extension_manager.get_provider().lock().await.clone();
        match provider {
            Some(provider) => provider
                .get_context_limit(model_config)
                .await
                .ok()
                .or_else(|| Some(model_config.context_limit())),
            None => Some(model_config.context_limit()),
        }
    } else {
        None
    };
    let total_tokens = session
        .as_ref()
        .and_then(|session| session.usage.total_tokens);
    let compaction_threshold = crate::config::Config::global()
        .get_param::<f64>("GOOSE_AUTO_COMPACT_THRESHOLD")
        .unwrap_or(crate::context_mgmt::DEFAULT_COMPACTION_THRESHOLD);
    compaction_remaining_line(total_tokens, context_limit, compaction_threshold)
}

/// The turn's context block: composed once per turn, persisted as an
/// agent-only user message, never moved or edited afterwards.
pub async fn turn_context_message(
    session_id: &str,
    extension_manager: &ExtensionManager,
    turns_taken: u32,
    max_turns: u32,
    turn_start: chrono::DateTime<chrono::Local>,
    compaction_info: Option<String>,
) -> Option<Message> {
    if SKIP.with(|f| f.get()) {
        return None;
    }

    let session = extension_manager
        .get_context()
        .session_manager
        .get_session(session_id, false)
        .await
        .ok();
    let session_model_config = session
        .as_ref()
        .and_then(|session| session.model_config.clone());
    let context_limit = if let Some(model_config) = session_model_config.as_ref() {
        let provider = extension_manager.get_provider().lock().await.clone();
        match provider {
            Some(provider) => provider
                .get_context_limit(model_config)
                .await
                .ok()
                .or_else(|| Some(model_config.context_limit())),
            None => Some(model_config.context_limit()),
        }
    } else {
        None
    };
    if should_skip_moim(context_limit) {
        return None;
    }

    let working_dir = session
        .as_ref()
        .map(|session| session.working_dir.clone())
        .unwrap_or_else(|| PathBuf::from("."));
    let mut parts = extension_manager.collect_moim_parts(session_id).await;
    parts.extend(compaction_info.map(|value| tag("compaction", &value)));
    parts.extend(turn_budget_part(turns_taken, max_turns));
    turn_context_event(&working_dir, context_limit, parts, turn_start)
}

/// The state-machine variant: operations contribute the parts and the caller
/// persists the event, so later requests in the turn reuse the same bytes.
pub(crate) fn turn_context_event(
    working_dir: &Path,
    context_limit: Option<usize>,
    parts: Vec<String>,
    turn_start: chrono::DateTime<chrono::Local>,
) -> Option<Message> {
    if SKIP.with(|f| f.get()) || should_skip_moim(context_limit) {
        return None;
    }

    Some(
        Message::user()
            .with_text(compose_moim(working_dir, parts, turn_start))
            .with_metadata(MessageMetadata::agent_only().with_turn_context()),
    )
}

fn should_skip_moim(context_limit: Option<usize>) -> bool {
    context_limit.is_some_and(|limit| limit < MIN_CONTEXT_FOR_MOIM)
}

fn compose_moim(
    working_dir: &Path,
    parts: Vec<String>,
    turn_start: chrono::DateTime<chrono::Local>,
) -> String {
    let timestamp = turn_start.format("%Y-%m-%d %H:%M:00 %:z");
    let mut lines = vec![
        open_tag(TURN_CONTEXT_TAG),
        tag(CURRENT_TIME_TAG, &timestamp.to_string()),
        tag(WORKING_DIRECTORY_TAG, &working_dir.display().to_string()),
    ];

    for part in parts {
        if !part.trim().is_empty() {
            lines.push(String::new());
            lines.push(part);
        }
    }

    lines.push(close_tag(TURN_CONTEXT_TAG));
    lines.join("\n")
}

fn open_tag(name: &str) -> String {
    format!("<{name}>")
}

fn close_tag(name: &str) -> String {
    format!("</{name}>")
}

fn tag(name: &str, value: &str) -> String {
    format!("<{name}>{}</{name}>", escape_xml_text(value))
}

fn escape_xml_text(value: &str) -> String {
    value
        .replace('&', "&amp;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
}

fn compaction_remaining_line(
    total_tokens: Option<i32>,
    context_limit: Option<usize>,
    threshold: f64,
) -> Option<String> {
    let total_tokens = total_tokens?;
    let context_limit = context_limit?;

    if total_tokens <= 0 || context_limit == 0 || threshold <= 0.0 || threshold >= 1.0 {
        return None;
    }

    let compaction_at = (context_limit as f64 * threshold) as i32;
    if compaction_at <= 0 || (total_tokens as f64 / compaction_at as f64) < 0.5 {
        return None;
    }

    Some(format!(
        "~{}k tokens remaining",
        compaction_at.saturating_sub(total_tokens) / 1000
    ))
}

fn turn_budget_part(turns_taken: u32, max_turns: u32) -> Option<String> {
    if max_turns == 0 || turns_taken.saturating_mul(2) < max_turns {
        return None;
    }

    Some(tag(
        "turn-budget",
        &format!("{turns_taken}/{max_turns} used"),
    ))
}

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

    async fn session_and_manager() -> (String, ExtensionManager, tempfile::TempDir) {
        let temp_dir = tempfile::tempdir().unwrap();
        let em = ExtensionManager::new_without_provider(temp_dir.path().to_path_buf());
        let session = em
            .get_context()
            .session_manager
            .create_session(
                PathBuf::from("/test/dir"),
                "test".to_string(),
                crate::session::SessionType::User,
                crate::config::GooseMode::Auto,
            )
            .await
            .unwrap();
        (session.id, em, temp_dir)
    }

    #[tokio::test]
    async fn turn_context_message_is_an_agent_only_user_message() {
        let (session_id, em, _tmp) = session_and_manager().await;

        let message = turn_context_message(&session_id, &em, 0, 100, chrono::Local::now(), None)
            .await
            .expect("turn context should be produced");

        assert_eq!(message.role, rmcp::model::Role::User);
        assert!(message.is_agent_visible());
        assert!(!message.is_user_visible());
        assert!(message.is_turn_context());
        let text = message.content[0].as_text().unwrap();
        assert!(text.starts_with(&format!("<{TURN_CONTEXT_TAG}>\n")));
        assert!(text.trim_end().ends_with(&format!("</{TURN_CONTEXT_TAG}>")));
    }

    #[tokio::test]
    async fn turn_context_bytes_are_stable_for_a_turn() {
        let (session_id, em, _tmp) = session_and_manager().await;
        let turn_start = chrono::Local::now();

        let first = turn_context_message(&session_id, &em, 0, 100, turn_start, None)
            .await
            .unwrap();
        let second = turn_context_message(&session_id, &em, 0, 100, turn_start, None)
            .await
            .unwrap();

        assert_eq!(
            first.content[0].as_text(),
            second.content[0].as_text(),
            "the same turn inputs must render byte-identical blocks"
        );
    }

    #[test]
    fn turn_context_event_is_marked_and_skips_small_contexts() {
        let event = turn_context_event(
            Path::new("/test/dir"),
            Some(200_000),
            vec![],
            chrono::Local::now(),
        )
        .expect("large-context models should get a turn-context event");
        assert!(event.is_turn_context());
        assert!(!event.is_user_visible());

        assert!(turn_context_event(
            Path::new("/test/dir"),
            Some(16_000),
            vec![],
            chrono::Local::now(),
        )
        .is_none());
    }

    #[test]
    fn current_time_is_timezone_aware() {
        let block = compose_moim(
            Path::new("/Users/me/code/goose"),
            vec![],
            chrono::Local::now(),
        );
        let time_line = block
            .lines()
            .find(|line| line.starts_with(&format!("<{CURRENT_TIME_TAG}>")))
            .expect("turn-context block should contain a current-time line");

        let value = time_line
            .trim_start_matches(&format!("<{CURRENT_TIME_TAG}>"))
            .trim_end_matches(&format!("</{CURRENT_TIME_TAG}>"));

        chrono::DateTime::parse_from_str(value, "%Y-%m-%d %H:%M:%S %:z")
            .expect("current-time should include a numeric UTC offset");
    }
}