llm-worker 0.2.0

//! Worker状態管理のテスト
//!
//! Type-stateパターン（Mutable/CacheLocked）による状態遷移と
//! ターン間の状態保持をテストする。

mod common;

use common::MockLlmClient;
use llm_worker::Worker;
use llm_worker::llm_client::event::{Event, ResponseStatus, StatusEvent};
use llm_worker::{Message, MessageContent};

// =============================================================================
// Mutable状態のテスト
// =============================================================================

/// Mutable状態でシステムプロンプトを設定できることを確認
#[test]
fn test_mutable_set_system_prompt() {
    let client = MockLlmClient::new(vec![]);
    let mut worker = Worker::new(client);

    assert!(worker.get_system_prompt().is_none());

    worker.set_system_prompt("You are a helpful assistant.");
    assert_eq!(
        worker.get_system_prompt(),
        Some("You are a helpful assistant.")
    );
}

/// Mutable状態で履歴を自由に編集できることを確認
#[test]
fn test_mutable_history_manipulation() {
    let client = MockLlmClient::new(vec![]);
    let mut worker = Worker::new(client);

    // 初期状態は空
    assert!(worker.history().is_empty());

    // 履歴を追加
    worker.push_message(Message::user("Hello"));
    worker.push_message(Message::assistant("Hi there!"));
    assert_eq!(worker.history().len(), 2);

    // 履歴への可変アクセス
    worker.history_mut().push(Message::user("How are you?"));
    assert_eq!(worker.history().len(), 3);

    // 履歴をクリア
    worker.clear_history();
    assert!(worker.history().is_empty());

    // 履歴を設定
    let messages = vec![Message::user("Test"), Message::assistant("Response")];
    worker.set_history(messages);
    assert_eq!(worker.history().len(), 2);
}

/// ビルダーパターンでWorkerを構築できることを確認
#[test]
fn test_mutable_builder_pattern() {
    let client = MockLlmClient::new(vec![]);
    let worker = Worker::new(client)
        .system_prompt("System prompt")
        .with_message(Message::user("Hello"))
        .with_message(Message::assistant("Hi!"))
        .with_messages(vec![
            Message::user("How are you?"),
            Message::assistant("I'm fine!"),
        ]);

    assert_eq!(worker.get_system_prompt(), Some("System prompt"));
    assert_eq!(worker.history().len(), 4);
}

/// extend_historyで複数メッセージを追加できることを確認
#[test]
fn test_mutable_extend_history() {
    let client = MockLlmClient::new(vec![]);
    let mut worker = Worker::new(client);

    worker.push_message(Message::user("First"));

    worker.extend_history(vec![
        Message::assistant("Response 1"),
        Message::user("Second"),
        Message::assistant("Response 2"),
    ]);

    assert_eq!(worker.history().len(), 4);
}

// =============================================================================
// 状態遷移テスト
// =============================================================================

/// lock()でMutable -> CacheLocked状態に遷移することを確認
#[test]
fn test_lock_transition() {
    let client = MockLlmClient::new(vec![]);
    let mut worker = Worker::new(client);

    worker.set_system_prompt("System");
    worker.push_message(Message::user("Hello"));
    worker.push_message(Message::assistant("Hi"));

    // ロック
    let locked_worker = worker.lock();

    // CacheLocked状態でも履歴とシステムプロンプトにアクセス可能
    assert_eq!(locked_worker.get_system_prompt(), Some("System"));
    assert_eq!(locked_worker.history().len(), 2);
    assert_eq!(locked_worker.locked_prefix_len(), 2);
}

/// unlock()でCacheLocked -> Mutable状態に遷移することを確認
#[test]
fn test_unlock_transition() {
    let client = MockLlmClient::new(vec![]);
    let mut worker = Worker::new(client);

    worker.push_message(Message::user("Hello"));
    let locked_worker = worker.lock();

    // アンロック
    let mut worker = locked_worker.unlock();

    // Mutable状態に戻ったので履歴操作が可能
    worker.push_message(Message::assistant("Hi"));
    worker.clear_history();
    assert!(worker.history().is_empty());
}

// =============================================================================
// ターン実行と状態保持のテスト
// =============================================================================

/// Mutable状態でターンを実行し、履歴が正しく更新されることを確認
#[tokio::test]
async fn test_mutable_run_updates_history() {
    let events = vec![
        Event::text_block_start(0),
        Event::text_delta(0, "Hello, I'm an assistant!"),
        Event::text_block_stop(0, None),
        Event::Status(StatusEvent {
            status: ResponseStatus::Completed,
        }),
    ];

    let client = MockLlmClient::new(events);
    let mut worker = Worker::new(client);

    // 実行
    let result = worker.run("Hi there").await;
    assert!(result.is_ok());

    // 履歴が更新されている
    let history = worker.history();
    assert_eq!(history.len(), 2); // user + assistant

    // ユーザーメッセージ
    assert!(matches!(
        &history[0].content,
        MessageContent::Text(t) if t == "Hi there"
    ));

    // アシスタントメッセージ
    assert!(matches!(
        &history[1].content,
        MessageContent::Text(t) if t == "Hello, I'm an assistant!"
    ));
}

/// CacheLocked状態で複数ターンを実行し、履歴が正しく累積することを確認
#[tokio::test]
async fn test_locked_multi_turn_history_accumulation() {
    // 2回のリクエストに対応するレスポンスを準備
    let client = MockLlmClient::with_responses(vec![
        // 1回目のレスポンス
        vec![
            Event::text_block_start(0),
            Event::text_delta(0, "Nice to meet you!"),
            Event::text_block_stop(0, None),
            Event::Status(StatusEvent {
                status: ResponseStatus::Completed,
            }),
        ],
        // 2回目のレスポンス
        vec![
            Event::text_block_start(0),
            Event::text_delta(0, "I can help with that."),
            Event::text_block_stop(0, None),
            Event::Status(StatusEvent {
                status: ResponseStatus::Completed,
            }),
        ],
    ]);

    let worker = Worker::new(client).system_prompt("You are helpful.");

    // ロック（システムプロンプト設定後）
    let mut locked_worker = worker.lock();
    assert_eq!(locked_worker.locked_prefix_len(), 0); // メッセージはまだない

    // 1ターン目
    let result1 = locked_worker.run("Hello!").await;
    assert!(result1.is_ok());
    assert_eq!(locked_worker.history().len(), 2); // user + assistant

    // 2ターン目
    let result2 = locked_worker.run("Can you help me?").await;
    assert!(result2.is_ok());
    assert_eq!(locked_worker.history().len(), 4); // 2 * (user + assistant)

    // 履歴の内容を確認
    let history = locked_worker.history();

    // 1ターン目のユーザーメッセージ
    assert!(matches!(&history[0].content, MessageContent::Text(t) if t == "Hello!"));

    // 1ターン目のアシスタントメッセージ
    assert!(matches!(&history[1].content, MessageContent::Text(t) if t == "Nice to meet you!"));

    // 2ターン目のユーザーメッセージ
    assert!(matches!(&history[2].content, MessageContent::Text(t) if t == "Can you help me?"));

    // 2ターン目のアシスタントメッセージ
    assert!(matches!(&history[3].content, MessageContent::Text(t) if t == "I can help with that."));
}

/// locked_prefix_lenがロック時点の履歴長を正しく記録することを確認
#[tokio::test]
async fn test_locked_prefix_len_tracking() {
    let client = MockLlmClient::with_responses(vec![
        vec![
            Event::text_block_start(0),
            Event::text_delta(0, "Response 1"),
            Event::text_block_stop(0, None),
            Event::Status(StatusEvent {
                status: ResponseStatus::Completed,
            }),
        ],
        vec![
            Event::text_block_start(0),
            Event::text_delta(0, "Response 2"),
            Event::text_block_stop(0, None),
            Event::Status(StatusEvent {
                status: ResponseStatus::Completed,
            }),
        ],
    ]);

    let mut worker = Worker::new(client);

    // 事前にメッセージを追加
    worker.push_message(Message::user("Pre-existing message 1"));
    worker.push_message(Message::assistant("Pre-existing response 1"));

    assert_eq!(worker.history().len(), 2);

    // ロック
    let mut locked_worker = worker.lock();
    assert_eq!(locked_worker.locked_prefix_len(), 2); // ロック時点で2メッセージ

    // ターン実行
    locked_worker.run("New message").await.unwrap();

    // 履歴は増えるが、locked_prefix_lenは変わらない
    assert_eq!(locked_worker.history().len(), 4); // 2 + 2
    assert_eq!(locked_worker.locked_prefix_len(), 2); // 変わらない
}

/// ターンカウントが正しくインクリメントされることを確認
#[tokio::test]
async fn test_turn_count_increment() {
    let client = MockLlmClient::with_responses(vec![
        vec![
            Event::text_block_start(0),
            Event::text_delta(0, "Turn 1"),
            Event::text_block_stop(0, None),
            Event::Status(StatusEvent {
                status: ResponseStatus::Completed,
            }),
        ],
        vec![
            Event::text_block_start(0),
            Event::text_delta(0, "Turn 2"),
            Event::text_block_stop(0, None),
            Event::Status(StatusEvent {
                status: ResponseStatus::Completed,
            }),
        ],
    ]);

    let mut worker = Worker::new(client);

    assert_eq!(worker.turn_count(), 0);

    worker.run("First").await.unwrap();
    assert_eq!(worker.turn_count(), 1);

    worker.run("Second").await.unwrap();
    assert_eq!(worker.turn_count(), 2);
}

/// unlock後に履歴を編集し、再度lockできることを確認
#[tokio::test]
async fn test_unlock_edit_relock() {
    let client = MockLlmClient::with_responses(vec![vec![
        Event::text_block_start(0),
        Event::text_delta(0, "Response"),
        Event::text_block_stop(0, None),
        Event::Status(StatusEvent {
            status: ResponseStatus::Completed,
        }),
    ]]);

    let worker = Worker::new(client)
        .with_message(Message::user("Hello"))
        .with_message(Message::assistant("Hi"));

    // ロック -> アンロック
    let locked = worker.lock();
    assert_eq!(locked.locked_prefix_len(), 2);

    let mut unlocked = locked.unlock();

    // 履歴を編集
    unlocked.clear_history();
    unlocked.push_message(Message::user("Fresh start"));

    // 再ロック
    let relocked = unlocked.lock();
    assert_eq!(relocked.history().len(), 1);
    assert_eq!(relocked.locked_prefix_len(), 1);
}

// =============================================================================
// システムプロンプト保持のテスト
// =============================================================================

/// CacheLocked状態でもシステムプロンプトが保持されることを確認
#[test]
fn test_system_prompt_preserved_in_locked_state() {
    let client = MockLlmClient::new(vec![]);
    let worker = Worker::new(client).system_prompt("Important system prompt");

    let locked = worker.lock();
    assert_eq!(locked.get_system_prompt(), Some("Important system prompt"));

    let unlocked = locked.unlock();
    assert_eq!(
        unlocked.get_system_prompt(),
        Some("Important system prompt")
    );
}

/// unlock -> 再lock でシステムプロンプトを変更できることを確認
#[test]
fn test_system_prompt_change_after_unlock() {
    let client = MockLlmClient::new(vec![]);
    let worker = Worker::new(client).system_prompt("Original prompt");

    let locked = worker.lock();
    let mut unlocked = locked.unlock();

    unlocked.set_system_prompt("New prompt");
    assert_eq!(unlocked.get_system_prompt(), Some("New prompt"));

    let relocked = unlocked.lock();
    assert_eq!(relocked.get_system_prompt(), Some("New prompt"));
}