magi-code 0.64.0

use super::*;

#[test]
fn context_usage_event_stores_reasoning_tokens_when_present() {
    let mut state = MissionControlState::default();
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 12,
        max_tokens: 128,
        reasoning_tokens: Some(34),
        source: ContextUsageSource::ProviderExact,
        request_sequence: 0,
    });

    assert_eq!(
        state.context_usage,
        Some(ContextUsageState {
            current_tokens: 12,
            max_tokens: 128,
            reasoning_tokens: Some(34),
            source: ContextUsageSource::ProviderExact,
            request_sequence: 0,
        })
    );
}

#[test]
fn context_usage_preserves_nonzero_count_after_tool_result_before_continuation_usage() {
    let call = crate::providers::ToolCall {
        id: "call_read".to_string(),
        name: "read".to_string(),
        arguments: serde_json::json!({"path":"src/lib.rs"}),
    };
    let result = crate::tools::ToolResult {
        tool_name: "read".to_string(),
        success: true,
        content: "contents".to_string(),
        metadata: serde_json::json!({}),
        display: crate::tools::ToolResultDisplay::default(),
    };
    let summary = crate::output::tool_display_summary(&call, &result);
    let mut state = MissionControlState {
        context_usage: Some(ContextUsageState {
            current_tokens: 12345,
            max_tokens: 128000,
            reasoning_tokens: None,
            source: ContextUsageSource::FallbackEstimate,
            request_sequence: 0,
        }),
        ..Default::default()
    };

    state.apply_activity_event(ActivityEvent::Started {
        id: ActivityId::new("call_read"),
        parent_id: None,
        kind: ActivityKind::Tool,
        status: ActivityStatus::Running,
        metadata: ActivityMetadata::new("read src/lib.rs"),
    });
    state.apply_activity_event(ActivityEvent::Finished {
        id: ActivityId::new("call_read"),
        status: ActivityStatus::Success,
        metadata: Some(ActivityMetadata::new("read src/lib.rs")),
    });
    state.apply_output_event(&OutputEvent::ToolResult {
        call: Box::new(call),
        result: Box::new(result),
        summary: Box::new(summary),
    });
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 0,
        max_tokens: 128000,
        reasoning_tokens: None,
        source: ContextUsageSource::FallbackEstimate,
        request_sequence: 0,
    });

    assert_eq!(
        state.context_usage,
        Some(ContextUsageState {
            current_tokens: 12345,
            max_tokens: 128000,
            reasoning_tokens: None,
            source: ContextUsageSource::FallbackEstimate,
            request_sequence: 0,
        })
    );

    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 13000,
        max_tokens: 128000,
        reasoning_tokens: None,
        source: ContextUsageSource::FallbackEstimate,
        request_sequence: 0,
    });

    assert_eq!(
        state.context_usage,
        Some(ContextUsageState {
            current_tokens: 13000,
            max_tokens: 128000,
            reasoning_tokens: None,
            source: ContextUsageSource::FallbackEstimate,
            request_sequence: 0,
        })
    );
}

#[test]
fn context_usage_accepts_initial_zero_without_prior_estimate() {
    let mut state = MissionControlState::default();
    state.apply_activity_event(ActivityEvent::Started {
        id: ActivityId::new("tool-1"),
        parent_id: None,
        kind: ActivityKind::Tool,
        status: ActivityStatus::Running,
        metadata: ActivityMetadata::new("read"),
    });

    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 0,
        max_tokens: 128000,
        reasoning_tokens: None,
        source: ContextUsageSource::FallbackEstimate,
        request_sequence: 0,
    });

    assert_eq!(
        state.context_usage,
        Some(ContextUsageState {
            current_tokens: 0,
            max_tokens: 128000,
            reasoning_tokens: None,
            source: ContextUsageSource::FallbackEstimate,
            request_sequence: 0,
        })
    );
}

#[test]
fn context_usage_exact_provider_replaces_tokenizer_projection() {
    let mut state = MissionControlState::default();
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 120,
        max_tokens: 1000,
        reasoning_tokens: None,
        source: ContextUsageSource::TokenizerProjection,
        request_sequence: 7,
    });
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 100,
        max_tokens: 1000,
        reasoning_tokens: Some(5),
        source: ContextUsageSource::ProviderExact,
        request_sequence: 7,
    });

    assert_eq!(
        state.context_usage,
        Some(ContextUsageState {
            current_tokens: 100,
            max_tokens: 1000,
            reasoning_tokens: Some(5),
            source: ContextUsageSource::ProviderExact,
            request_sequence: 7,
        })
    );
}

#[test]
fn context_usage_partial_zero_preserves_known_nonzero_count() {
    let mut state = MissionControlState::default();
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 120,
        max_tokens: 1000,
        reasoning_tokens: None,
        source: ContextUsageSource::TokenizerProjection,
        request_sequence: 1,
    });
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 0,
        max_tokens: 1000,
        reasoning_tokens: Some(9),
        source: ContextUsageSource::ProviderPartial,
        request_sequence: 1,
    });

    assert_eq!(
        state.context_usage,
        Some(ContextUsageState {
            current_tokens: 120,
            max_tokens: 1000,
            reasoning_tokens: Some(9),
            source: ContextUsageSource::TokenizerProjection,
            request_sequence: 1,
        })
    );
}

#[test]
fn context_usage_fallback_same_sequence_does_not_replace_exact() {
    let mut state = MissionControlState::default();
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 100,
        max_tokens: 1000,
        reasoning_tokens: None,
        source: ContextUsageSource::ProviderExact,
        request_sequence: 1,
    });
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 130,
        max_tokens: 1000,
        reasoning_tokens: None,
        source: ContextUsageSource::FallbackEstimate,
        request_sequence: 1,
    });

    assert_eq!(state.context_usage.unwrap().current_tokens, 100);
    assert_eq!(
        state.context_usage.unwrap().source,
        ContextUsageSource::ProviderExact
    );
}

#[test]
fn context_usage_resets_freshness_for_new_prompt_run() {
    let mut state = MissionControlState::default();
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 500,
        max_tokens: 1000,
        reasoning_tokens: Some(9),
        source: ContextUsageSource::ProviderExact,
        request_sequence: 1,
    });
    state.apply_output_event(&OutputEvent::UserPrompt {
        text: "next prompt".to_string(),
    });
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 700,
        max_tokens: 1000,
        reasoning_tokens: None,
        source: ContextUsageSource::TokenizerEstimate,
        request_sequence: 0,
    });

    assert_eq!(
        state.context_usage,
        Some(ContextUsageState {
            current_tokens: 700,
            max_tokens: 1000,
            reasoning_tokens: None,
            source: ContextUsageSource::TokenizerEstimate,
            request_sequence: 0,
        })
    );
}

#[test]
fn context_usage_new_sequence_estimate_can_replace_previous_exact() {
    let mut state = MissionControlState::default();
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 100,
        max_tokens: 1000,
        reasoning_tokens: Some(3),
        source: ContextUsageSource::ProviderExact,
        request_sequence: 1,
    });
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 150,
        max_tokens: 1000,
        reasoning_tokens: None,
        source: ContextUsageSource::TokenizerEstimate,
        request_sequence: 2,
    });

    assert_eq!(
        state.context_usage,
        Some(ContextUsageState {
            current_tokens: 150,
            max_tokens: 1000,
            reasoning_tokens: Some(3),
            source: ContextUsageSource::TokenizerEstimate,
            request_sequence: 2,
        })
    );
}

#[test]
fn context_usage_reasoning_tokens_remain_optional() {
    let mut state = MissionControlState::default();
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 100,
        max_tokens: 1000,
        reasoning_tokens: Some(3),
        source: ContextUsageSource::ProviderExact,
        request_sequence: 1,
    });
    state.apply_output_event(&OutputEvent::ContextUsage {
        current_tokens: 101,
        max_tokens: 1000,
        reasoning_tokens: None,
        source: ContextUsageSource::ProviderExact,
        request_sequence: 2,
    });

    assert_eq!(state.context_usage.unwrap().reasoning_tokens, None);
}