harn-vm 0.10.167

use super::*;
use std::sync::atomic::{AtomicUsize, Ordering};

use crate::agent_events::{
    emit_event, register_sink, AgentEvent, AgentEventSink, AgentEventTransport,
};
use crate::value::VmDictExt;

struct Count(Arc<AtomicUsize>);
impl AgentEventSink for Count {
    fn handle_event(&self, _: &AgentEvent) {
        self.0.fetch_add(1, Ordering::SeqCst);
    }
}

fn transport() -> (AgentEventTransport, Arc<AtomicUsize>) {
    let count = Arc::new(AtomicUsize::new(0));
    (
        AgentEventTransport::new(Arc::new(Count(count.clone()))),
        count,
    )
}

fn event(session: &str) -> AgentEvent {
    AgentEvent::IterationStart {
        session_id: session.into(),
        iteration: 1,
        provider: String::new(),
        model: String::new(),
    }
}

#[tokio::test(flavor = "current_thread")]
async fn never_executed_vm_cleanup_keeps_transport_explicitly_empty() {
    let (ambient, count) = transport();
    ambient
        .scope(async {
            let vm = crate::Vm::new();
            emit_event(&event("never-executed-cleanup"));
            assert_eq!(count.load(Ordering::SeqCst), 1);
            ScopedCleanup {
                runtimes: vm.agent_cleanup_runtimes(),
                inner: async {
                    emit_event(&event("never-executed-cleanup"));
                },
            }
            .await;
            assert_eq!(count.load(Ordering::SeqCst), 1);
            emit_event(&event("never-executed-cleanup"));
            assert_eq!(count.load(Ordering::SeqCst), 2);
        })
        .await;
}

#[tokio::test(flavor = "current_thread")]
async fn reused_vm_captures_current_runtimes_and_retains_them_for_children_and_cleanup() {
    let mut vm = crate::Vm::new();
    crate::register_vm_stdlib(&mut vm);
    let chunk = crate::compile_source("pipeline main() {}").unwrap();
    let (old, old_count) = transport();
    let (new, new_count) = transport();
    let first = old
        .scope(async {
            vm.execute(&chunk).await.unwrap();
            vm.agent_cleanup_runtimes()
        })
        .await;
    vm.session_runtime = crate::agent_sessions::fresh_session_runtime();
    vm.agent_host_session_runtime =
        crate::llm::agent_session_host::fresh_agent_host_session_runtime();
    let second = new
        .scope(async {
            vm.execute(&chunk).await.unwrap();
            vm.agent_cleanup_runtimes()
        })
        .await;
    let first_key = first.key();
    let second_key = second.key();
    assert_ne!(first_key.execution_id, second_key.execution_id);
    assert_ne!(first_key.session_runtime, second_key.session_runtime);
    assert_ne!(first_key.host_runtime, second_key.host_runtime);
    assert_eq!(
        second_key.session_runtime,
        Arc::as_ptr(&vm.session_runtime) as usize
    );
    assert_eq!(
        second_key.host_runtime,
        Arc::as_ptr(&vm.agent_host_session_runtime) as usize
    );
    assert_eq!(vm.child_vm().agent_cleanup_runtimes().key(), second_key);

    let retained = second.clone();
    assert_eq!(retained.key(), second_key);
    let old_before = old_count.load(Ordering::SeqCst);
    let new_before = new_count.load(Ordering::SeqCst);
    // Scheduling under OLD ambient state must use the complete retained NEW
    // snapshot, including its unchanged execution identity.
    old.scope(ScopedCleanup {
        runtimes: retained,
        inner: async {
            emit_event(&event("reused-vm-cleanup"));
        },
    })
    .await;
    assert_eq!(old_count.load(Ordering::SeqCst), old_before);
    assert_eq!(new_count.load(Ordering::SeqCst), new_before + 1);
}

#[tokio::test(flavor = "current_thread")]
async fn top_level_vm_drop_keeps_origin_after_execution_scope_has_unwound() {
    let root = tempfile::tempdir().unwrap();
    let session = format!("transport-drop-{}", uuid::Uuid::now_v7());
    let (old, old_count) = transport();
    let (new, new_count) = transport();
    let vm = old
        .scope(async {
            let mut vm = crate::Vm::new();
            crate::register_vm_stdlib(&mut vm);
            let chunk = crate::compile_source("pipeline main() {}").unwrap();
            vm.execute(&chunk).await.unwrap();
            vm
        })
        .await;
    let old_before_cleanup = old_count.load(Ordering::SeqCst);
    let mut options = crate::value::DictMap::new();
    options.put_str("root", root.path().to_string_lossy().as_ref());
    let prepared = crate::agent_session_journal::prepare(
        &session,
        &options,
        format!("run-{session}"),
        format!("turn-{session}"),
    )
    .await
    .unwrap();
    crate::agent_sessions::open_or_create_for_test(Some(session.clone()));
    crate::agent_sessions::install_journal(&session, prepared.state).unwrap();
    crate::agent_sessions::claim_journal_task(
        &session,
        vm.execution_id(),
        "task_root".into(),
        true,
    )
    .unwrap();
    crate::llm::agent_session_host::seed_host_session_provider_model(&session, "mock", "fixture");
    let durable = Arc::new(AtomicUsize::new(0));
    register_sink(session.clone(), Arc::new(Count(durable.clone())));
    assert_eq!(
        crate::agent_events::session_external_sink_count(&session),
        1,
        "the observer census reaches a known nonzero registration"
    );
    let mut progress = crate::agent_lifecycle_cleanup::subscribe_cleanup_progress();
    let before = crate::vm::subtask::lifecycle_cleanup_spawn_count();
    new.scope(async {
        emit_event(&event(&session));
        // Drop happens under NEW ambient ownership, after OLD execution unwound.
        drop(vm);
        assert!(crate::vm::subtask::lifecycle_cleanup_spawn_count() > before);
        crate::agent_lifecycle_cleanup::settle_cleanup(
            &mut progress,
            || {
                !crate::agent_sessions::has_journal(&session)
                    && !crate::agent_sessions::exists(&session)
                    && crate::agent_events::session_external_sink_count(&session) == 0
                    && old_count.load(Ordering::SeqCst) > old_before_cleanup
            },
            "originating top-level cleanup must emit and release both session owners",
        )
        .await;
        let cleanup_events = old_count.load(Ordering::SeqCst) - old_before_cleanup;
        assert!(
            cleanup_events > 0,
            "actual terminal cleanup reaches OLD transport"
        );
        let observed_at_close = durable.load(Ordering::SeqCst);
        assert_eq!(
            observed_at_close,
            cleanup_events + 1,
            "session observer receives every cleanup event and the pre-close current anchor"
        );
        assert_eq!(
            crate::agent_events::session_external_sink_count(&session),
            0,
            "settled session close releases its dynamic observers"
        );
        // This post-close positive proves the caller's NEW transport was
        // restored; a closed session no longer owns a registry subscription.
        emit_event(&event(&session));
        assert_eq!(
            durable.load(Ordering::SeqCst),
            observed_at_close,
            "post-close transport observation cannot revive session subscribers"
        );
    })
    .await;
    crate::agent_events::clear_session_sinks(&session);
    assert_eq!(
        new_count.load(Ordering::SeqCst),
        2,
        "cleanup never uses NEW transport"
    );
    // The Count fixture measures observation, not durability. Read the owning
    // journal independently after close to prove its terminal really committed.
    let store = crate::stdlib::session_store::open_canonical_agent_session(
        &crate::stdlib::session_store::SessionStoreDir::under_root(root.path()),
        &session,
        None,
        harn_session_store::SessionType::User,
    )
    .await
    .expect("open canonical session after cleanup");
    let events = crate::stdlib::session_store::read_all_events(&store, &session)
        .await
        .expect("read committed journal after cleanup");
    let terminals: Vec<_> = events
        .iter()
        .filter(|event| {
            event
                .payload
                .pointer("/transcript_event/kind")
                .and_then(serde_json::Value::as_str)
                == Some("agent_run_terminal")
        })
        .collect();
    assert_eq!(
        terminals.len(),
        1,
        "cleanup commits exactly one durable terminal"
    );
    let metadata = terminals[0]
        .payload
        .pointer("/transcript_event/metadata")
        .unwrap();
    assert_eq!(metadata["final_status"], "cancelled");
    assert_eq!(metadata["stop_reason"], "cancelled");
    let phase = crate::agent_events::AgentTurnPhase::from_terminal_record(metadata)
        .expect("persisted terminal decodes through owning typed projection");
    assert!(
        matches!(phase, crate::agent_events::AgentTurnPhase::Terminal { outcome, .. }
        if outcome.kind == crate::agent_events::AgentTerminalKind::UserCancelled)
    );
}

async fn pending_cleanup_after_vm_reuse(drop_instead_of_retry: bool) {
    let root = tempfile::tempdir().expect("temp root");
    let session = format!("pending-reuse-{}", uuid::Uuid::now_v7());
    let task_id = "task_pending_reuse";
    let public_id = "public-pending-reuse";
    let (old, old_count) = transport();
    let (new, new_count) = transport();
    let mut vm = crate::Vm::new();
    crate::register_vm_stdlib(&mut vm);
    let chunk = crate::compile_source("pipeline main() {}").unwrap();
    old.scope(vm.execute(&chunk)).await.unwrap();
    let original = vm.agent_cleanup_runtimes();
    let original_execution = vm.execution_id().to_string();
    let store = ScopedCleanup {
        runtimes: original.clone(),
        inner: async {
            let mut options = crate::value::DictMap::new();
            options.put_str("root", root.path().to_string_lossy().as_ref());
            let prepared = crate::agent_session_journal::prepare(
                &session,
                &options,
                format!("run-{session}"),
                format!("turn-{session}"),
            )
            .await
            .unwrap();
            let store = prepared.state.store();
            crate::agent_sessions::open_or_create_for_test(Some(session.clone()));
            crate::agent_sessions::install_journal(&session, prepared.state).unwrap();
            crate::agent_sessions::claim_journal_task(
                &session,
                &original_execution,
                task_id.into(),
                true,
            )
            .unwrap();
            // Fail a real canonical append without permanently closing the session.
            // This private database's removable trigger injects only the write fault.
            let fault = rusqlite::Connection::open(store.path()).unwrap();
            fault
                .execute_batch(
                    "CREATE TRIGGER fail_pending_cleanup BEFORE INSERT ON session_events
                 BEGIN SELECT RAISE(FAIL, 'pending cleanup append fault'); END;",
                )
                .unwrap();
            vm.register_spawned_task(
                public_id.into(),
                crate::value::VmTaskHandle {
                    handle: tokio::spawn(std::future::pending()),
                    cancel_token: Arc::new(std::sync::atomic::AtomicBool::new(false)),
                    wait_task_id: task_id.into(),
                },
            );
            let error = vm
                .cancel_task_for_test(public_id)
                .await
                .expect_err("actual cancel must observe the canonical append failure");
            assert!(error.to_string().contains("pending cleanup append fault"));
            assert!(!vm.spawned_tasks.contains_key(public_id));
            assert!(vm.pending_task_cleanups.contains_key(public_id));
            assert!(crate::agent_sessions::has_journal(&session));
            assert!(crate::agent_sessions::next_journal_event(&session)
                .unwrap()
                .is_some());
            fault
                .execute_batch("DROP TRIGGER fail_pending_cleanup;")
                .unwrap();
            store
        },
    }
    .await;

    let old_before = old_count.load(Ordering::SeqCst);
    vm.session_runtime = crate::agent_sessions::fresh_session_runtime();
    vm.agent_host_session_runtime =
        crate::llm::agent_session_host::fresh_agent_host_session_runtime();
    new.scope(vm.execute(&chunk)).await.unwrap();
    assert_ne!(vm.execution_id().to_string(), original_execution);
    assert!(vm.pending_task_cleanups.contains_key(public_id));
    let current = vm.agent_cleanup_runtimes();
    let new_before = new_count.load(Ordering::SeqCst);
    let mut progress = subscribe_cleanup_progress();
    ScopedCleanup {
        runtimes: current,
        inner: async {
            if drop_instead_of_retry {
                drop(vm);
            } else {
                assert!(vm.cancel_task_for_test(public_id).await
                    .expect("retry commits the original pending terminal"));
                assert!(!vm.pending_task_cleanups.contains_key(public_id));
                ScopedCleanup {
                    runtimes: original.clone(),
                    inner: async {
                        assert!(
                            !crate::agent_sessions::has_journal(&session),
                            "a successful retry must discharge the original journal, not find zero sessions in the new runtime"
                        );
                        assert!(old_count.load(Ordering::SeqCst) > old_before);
                    },
                }.await;
                drop(vm);
            }
            ScopedCleanup {
                runtimes: original,
                inner: async {
                    settle_cleanup(
                        &mut progress,
                        || {
                            !crate::agent_sessions::has_journal(&session)
                                && !crate::agent_sessions::exists(&session)
                                && old_count.load(Ordering::SeqCst) > old_before
                        },
                        "pending cleanup after reuse must commit and emit through its original owner",
                    ).await;
                },
            }.await;
            emit_event(&event(&session));
        },
    }.await;
    assert!(old_count.load(Ordering::SeqCst) > old_before);
    assert_eq!(new_count.load(Ordering::SeqCst), new_before + 1);
    let events = crate::stdlib::session_store::read_all_events(&store, &session)
        .await
        .unwrap();
    let terminals: Vec<_> = events
        .iter()
        .filter(|event| event.payload["transcript_event"]["kind"] == "agent_run_terminal")
        .collect();
    assert_eq!(terminals.len(), 1);
    assert_eq!(
        terminals[0].payload["transcript_event"]["metadata"]["final_status"],
        "cancelled"
    );
    assert_eq!(
        terminals[0].payload["transcript_event"]["metadata"]["stop_reason"],
        "cancelled"
    );
}

#[tokio::test(flavor = "current_thread")]
async fn pending_cancel_retry_after_vm_reuse_retains_original_owner() {
    pending_cleanup_after_vm_reuse(false).await;
}

#[tokio::test(flavor = "current_thread")]
async fn pending_cancel_drop_after_vm_reuse_retains_original_owner() {
    pending_cleanup_after_vm_reuse(true).await;
}

#[tokio::test(flavor = "current_thread")]
async fn active_child_cancel_after_same_runtime_vm_reuse_retains_original_owner() {
    let root = tempfile::tempdir().unwrap();
    let session = format!("active-reuse-{}", uuid::Uuid::now_v7());
    let task_id = "task_active_reuse";
    let public_id = "public-active-reuse";
    let (old, old_count) = transport();
    let (new, new_count) = transport();
    let mut vm = crate::Vm::new();
    crate::register_vm_stdlib(&mut vm);
    let chunk = crate::compile_source("pipeline main() {}").unwrap();
    old.scope(vm.execute(&chunk)).await.unwrap();
    let original = vm.agent_cleanup_runtimes();
    let original_key = original.key();
    let cancelled = Arc::new(std::sync::atomic::AtomicBool::new(false));
    let store = ScopedCleanup {
        runtimes: original.clone(),
        inner: async {
            let mut options = crate::value::DictMap::new();
            options.put_str("root", root.path().to_string_lossy().as_ref());
            let prepared = crate::agent_session_journal::prepare(
                &session,
                &options,
                format!("run-{session}"),
                format!("turn-{session}"),
            )
            .await
            .unwrap();
            let store = prepared.state.store();
            crate::agent_sessions::open_or_create_for_test(Some(session.clone()));
            crate::agent_sessions::install_journal(&session, prepared.state).unwrap();
            crate::agent_sessions::claim_journal_task(
                &session,
                &original_key.execution_id,
                task_id.into(),
                true,
            )
            .unwrap();
            let (started_tx, started_rx) = tokio::sync::oneshot::channel();
            vm.register_spawned_task(
                public_id.into(),
                crate::value::VmTaskHandle {
                    handle: tokio::spawn(async move {
                        started_tx.send(()).unwrap();
                        std::future::pending().await
                    }),
                    cancel_token: cancelled.clone(),
                    wait_task_id: task_id.into(),
                },
            );
            started_rx
                .await
                .expect("the original child actually started");
            assert!(crate::agent_sessions::has_journal(&session));
            assert!(!vm.spawned_tasks[public_id].task.handle.is_finished());
            store
        },
    }
    .await;
    let old_before = old_count.load(Ordering::SeqCst);
    // Keep BOTH runtime handles unchanged: only top-level execution identity
    // and transport advance, unlike the pending-write-failure controls.
    new.scope(vm.execute(&chunk)).await.unwrap();
    let current = vm.agent_cleanup_runtimes();
    let current_key = current.key();
    assert_ne!(current_key.execution_id, original_key.execution_id);
    assert_eq!(current_key.session_runtime, original_key.session_runtime);
    assert_eq!(current_key.host_runtime, original_key.host_runtime);
    assert!(vm.spawned_tasks.contains_key(public_id));
    assert!(!cancelled.load(Ordering::SeqCst));
    let new_before = new_count.load(Ordering::SeqCst);
    ScopedCleanup {
        runtimes: current,
        inner: async {
            emit_event(&event(&session));
            assert_eq!(new_count.load(Ordering::SeqCst), new_before + 1);
            assert!(vm.cancel_task_for_test(public_id).await.unwrap());
            assert!(cancelled.load(Ordering::SeqCst));
            assert!(!vm.spawned_tasks.contains_key(public_id));
            assert!(!vm.pending_task_cleanups.contains_key(public_id));
            ScopedCleanup {
                runtimes: original,
                inner: async {
                    assert!(
                        !crate::agent_sessions::has_journal(&session),
                        "cancelling an active old child after reuse must commit its original terminal"
                    );
                    assert!(!crate::agent_sessions::exists(&session));
                    assert!(old_count.load(Ordering::SeqCst) > old_before);
                },
            }.await;
            emit_event(&event(&session));
            assert_eq!(new_count.load(Ordering::SeqCst), new_before + 2);
        },
    }.await;
    let events = crate::stdlib::session_store::read_all_events(&store, &session)
        .await
        .unwrap();
    let terminals: Vec<_> = events
        .iter()
        .filter(|event| event.payload["transcript_event"]["kind"] == "agent_run_terminal")
        .collect();
    assert_eq!(terminals.len(), 1);
    assert_eq!(
        terminals[0].payload["transcript_event"]["metadata"]["final_status"],
        "cancelled"
    );
    assert_eq!(
        terminals[0].payload["transcript_event"]["metadata"]["stop_reason"],
        "cancelled"
    );
}