plugins/agent/native/src/owner/tests/todo.rs

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

use super::*;
use maka_agent_facts::{
    event::{Fact, ToolOutcome},
    tool_call::{ToolCallIdentity, ToolRejection},
    tools::{ToolCallContext, ToolJournal, ToolPreparer},
};

async fn fixture(base: &str) -> Fixture {
    Fixture::configured(base, None, None, None, |settings| {
        settings.todo = true;
        settings.permissions.access = WorkspaceAccess::ReadOnly;
    })
    .await
}
fn last_tool(sent: &Value) -> &str {
    let messages = sent["messages"].as_array().unwrap();
    let tool = messages
        .iter()
        .rev()
        .find(|message| message["role"] == "tool")
        .unwrap_or_else(|| {
            panic!(
                "Expected a tool result; received: {:?}",
                messages
                    .iter()
                    .map(|message| (
                        message["role"].clone(),
                        message["content"]
                            .to_string()
                            .chars()
                            .take(1000)
                            .collect::<String>()
                    ))
                    .collect::<Vec<_>>()
            )
        });
    tool["content"].as_str().unwrap()
}
async fn page(f: &Fixture, session: &str, offset: usize) -> Value {
    f.query("checklist", json!({"sessionId":session,"offset":offset}))
        .await
        .unwrap()
}
pub(super) async fn reopen(f: &mut Fixture) {
    let deployment = f.app.source().unwrap().deployment;
    f.app.shutdown().await.unwrap();
    let mut native = NativeRegistry::default();
    maka_settings_native::provide(&mut native, crate::owner::tests::bundled_catalog()).unwrap();
    provide_all(&mut native).unwrap();
    f.app = Application::start_from_store(
        f.store.clone(),
        deployment,
        native,
        f.issuer.clone(),
        Arc::new(Grant(f.issuer.clone(), None)),
        Some(maka_host_js::factory(Default::default())),
    )
    .await
    .unwrap();
    f.app.ready().await.unwrap();
}
#[tokio::test]
async fn todo_archived_result_next_is_readable_by_the_selected_native_read_tool() {
    tokio::time::timeout(Duration::from_secs(90), async {
        for reader in ["Read", "WorkspaceFile"] {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let mut setup = service_tools::Setup::files("readText").named(reader);
        setup.nestable = true;
        let original_description = setup.selection().description;
        let policy = setup.policy();
        let mut f = Fixture::configured(
            &format!("http://{}/v1", listener.local_addr().unwrap()),
            None,
            None,
            Some(setup),
            |settings| {
                settings.todo = true;
                settings.permissions.access = WorkspaceAccess::ReadOnly;
            },
        )
        .await;
        // Declare the controlled provider's capacity, keeping this recovery
        // proof independent of an unrelated automatic compaction conversation.
        selection::edit(&f, |changes| {
            changes.model_overrides = Patch::Set(BTreeMap::from([("fixture-model".into(), ModelOverride {
                context_window: Some(1_000_000), input_limit: Some(900_000), ..Default::default()
            })]));
        }).await;
        let session = f.create("Recover original checklist result").await;
        let other = f.create("Cannot read another checklist archive").await;
        let items: Vec<Value> = (0..32)
            .map(|index| json!({"content":format!("{index:02} {}", "中文😀".repeat(if index == 0 { 2048 } else { 128 })),"status":"pending"}))
            .collect();
        // Independent literal rendering oracle: no Document::render or pager.
        let expected = (0..32).map(|index| format!("{}. [pending] \"{index:02} {}\"", index + 1,
            "中文😀".repeat(if index == 0 { 2048 } else { 128 }))).collect::<Vec<_>>().join("\n");
        f.send(&session).await.unwrap();
        let (mut socket, _) = listener.accept().await.unwrap();
        let sent = request(&mut socket).await;
        let selected_reader = sent["tools"].as_array().unwrap().iter()
            .find(|tool| tool["function"]["name"] == reader).cloned();
        let advertises_read = sent["tools"].as_array().unwrap().iter()
            .any(|tool| tool["function"]["name"] == "Read");
        service_tools::named_tool_response(&mut socket, "todo_write", json!({"todos":items})).await;

        let (mut socket, _) = listener.accept().await.unwrap();
        let sent = request(&mut socket).await;
        let archived: Value = serde_json::from_str(last_tool(&sent)).unwrap();
        let first_next = archived["page"]["next"].clone();
        let mut pages = vec![archived["page"].clone()];
        let mut recovery_results = Vec::new();
        let mut next = first_next.clone();
        while next.is_object() {
            // Every page is requested through the actual model with exactly the
            // prior tool's next object, never a manufactured cursor/store read.
            service_tools::named_tool_response(&mut socket, reader, next).await;
            let (mut received, _) = listener.accept().await.unwrap();
            let sent = request(&mut received).await;
            let result: Value = serde_json::from_str(last_tool(&sent)).unwrap();
            next = result["value"]["next"].clone();
            if result["kind"] == "completed" {
                pages.push(result["value"].clone());
            }
            recovery_results.push(result);
            socket = received;
        }
        response(&mut socket).await;
        f.drained().await;

        f.send(&other).await.unwrap();
        let (mut socket, _) = listener.accept().await.unwrap();
        request(&mut socket).await;
        service_tools::named_tool_response(&mut socket, reader, first_next.clone()).await;
        let (mut socket, _) = listener.accept().await.unwrap();
        let sent = request(&mut socket).await;
        let foreign: Value = serde_json::from_str(last_tool(&sent)).unwrap();
        response(&mut socket).await;
        f.drained().await;

        let code_result = if reader == "WorkspaceFile" {
            selection::edit(&f, |changes| {
                changes.model_overrides = Patch::Set(BTreeMap::from([("fixture-model".into(), ModelOverride {
                    code_mode: Some(true), context_window: Some(1_000_000), input_limit: Some(900_000), ..Default::default()
                })]));
            }).await;
            f.send(&session).await.unwrap();
            let (mut socket, _) = listener.accept().await.unwrap();
            let code_request = request(&mut socket).await;
            let code = format!(
                "const reader = ALL_TOOLS.find(tool => tool.name === 'WorkspaceFile'); if (!reader || !reader.description.includes({description}) || !reader.description.includes('archived tool-result') || !reader.description.includes('page.next')) throw Error('archive reader metadata missing'); const actual = await tools.WorkspaceFile({next}); const expected = {expected}; if (JSON.stringify(actual) !== JSON.stringify(expected)) throw Error('archive page differs from Direct'); text('k68-alias-archive-recovered');",
                description = serde_json::to_string(&original_description).unwrap(),
                next = first_next,
                expected = recovery_results[0],
            );
            service_tools::named_tool_response(&mut socket, "exec", json!({"code":code})).await;
            let (mut socket, _) = listener.accept().await.unwrap();
            let code_response = request(&mut socket).await;
            let output = last_tool(&code_response).to_owned();
            response(&mut socket).await;
            f.drained().await;
            selection::edit(&f, |changes| {
                changes.model_overrides = Patch::Set(BTreeMap::from([("fixture-model".into(), ModelOverride {
                    code_mode: Some(false), context_window: Some(1_000_000), input_limit: Some(900_000), ..Default::default()
                })]));
            }).await;
            Some((code_request, output))
        } else { None };

        let deployment = f.app.source().unwrap().deployment;
        f.app.shutdown().await.unwrap();
        let mut native = NativeRegistry::default();
        maka_settings_native::provide(&mut native, crate::owner::tests::bundled_catalog()).unwrap();
        provide_all(&mut native).unwrap();
        f.app = Application::start_from_store(
            f.store.clone(),
            deployment,
            native,
            f.issuer.clone(),
            Arc::new(Grant(f.issuer.clone(), Some(policy))),
            Some(maka_host_js::factory(Default::default())),
        ).await.unwrap();
        f.app.ready().await.unwrap();
        // send uses the current Application's public service. The old fixture
        // owner is retired; final Application shutdown drains the new owner.
        f.send(&session).await.unwrap();
        let (mut socket, _) = listener.accept().await.unwrap();
        request(&mut socket).await;
        service_tools::named_tool_response(&mut socket, reader, first_next.clone()).await;
        let (mut socket, _) = listener.accept().await.unwrap();
        let sent = request(&mut socket).await;
        let restarted_raw = last_tool(&sent).to_owned();
        response(&mut socket).await;
        while !f.query("observe", json!({"sessionId":session})).await.unwrap()["run"].is_null() {
            tokio::task::yield_now().await;
        }
        f.app.shutdown().await.unwrap();

        let original = f.store.scoped_prefix(
            maka_agent_facts::event::LogScope::Session { id: session.clone() },
            256,
            4 * 1024 * 1024,
        ).await.unwrap();
        let source_operation = original.events.iter().find_map(|event| match &event.event.fact {
            Fact::ToolDispatched { operation_id, name, .. } if name == "todo_write" => Some(operation_id),
            _ => None,
        }).unwrap();
        let source = original.events.iter().find(|event| matches!(&event.event.fact,
            Fact::ToolSettled { operation_id, outcome: ToolOutcome::Succeeded { .. } }
            if operation_id == source_operation)).unwrap();
        let source_output = f.store.resolve_tool_result(&session, &source.event.id).await.unwrap();
        f.config.shutdown().await.unwrap();

        // All behavior assertions follow actual cleanup, including expected
        // failure paths, so no failing recovery assertion abandons the Host.
        let restarted: Value = serde_json::from_str(&restarted_raw)
            .unwrap_or_else(|error| panic!("restarted Read returned non-JSON after cleanup: {error}; raw={restarted_raw:?}"));
        let selected_reader = selected_reader.expect("the original model request must select the genuine reader binding");
        let description = selected_reader["function"]["description"].as_str().unwrap();
        assert!(description.contains(&original_description), "{selected_reader}");
        assert!(description.contains("archived tool-result") && description.contains("page.next"), "{selected_reader}");
        if let Some((code_request, output)) = code_result {
            assert!(!advertises_read, "an aliased reader must not invent a Read tool");
            assert!(code_request["tools"].as_array().unwrap().iter()
                .any(|tool| tool["function"]["name"] == "exec"));
            assert_eq!(output.matches("k68-alias-archive-recovered").count(), 1, "{output}");
            assert!(!output.contains("\"ok\":false"), "{output}");
        }
        assert_eq!(archived["kind"], "maka.archived_tool_result");
        assert!(archived.to_string().encode_utf16().count()
            <= maka_agent_facts::read::MAX_PAGE_CHARS);
        assert_eq!(archived["toolName"], "todo_write");
        assert!(first_next.is_object(), "the oversized original must advertise next");
        let maka_agent_facts::tool_output::ToolOutput::Json(document) = source_output else {
            panic!("the original checklist output must remain JSON");
        };
        assert_eq!(document["items"], json!(items));
        let Fact::ToolSettled { outcome: ToolOutcome::Succeeded { model_projection, .. }, .. } = &source.event.fact else { unreachable!() };
        assert_eq!(model_projection, &maka_agent_facts::tool_output::DurableToolProjection::Text { text: expected.clone() });
        assert_eq!(original.events.iter().filter(|event| matches!(&event.event.fact,
            Fact::ToolDispatched { name, .. } if name == "todo_write")).count(), 1);
        for result in &recovery_results {
            assert_eq!(result["kind"], "completed", "the selected Read must recover the actual archive continuation: {result}");
            assert!(serde_json::to_string(result).unwrap().encode_utf16().count()
                <= maka_agent_facts::read::MAX_PAGE_CHARS, "recovery pages must not be archived again");
        }
        let mut recovered = String::new();
        let mut omitted_newline = false;
        for (index, page) in pages.iter().enumerate() {
            let content = page["content"].as_str().unwrap();
            recovered.push_str(content);
            assert_eq!(page["totalLines"], 32);
            if let Some(next) = pages.get(index + 1) {
                let end_line = page["offset"].as_u64().unwrap() + content.bytes().filter(|byte| *byte == b'\n').count() as u64;
                let next_line = next["offset"].as_u64().unwrap();
                if next_line == end_line + 1 {
                    recovered.push('\n');
                    omitted_newline = true;
                } else {
                    assert_eq!(next_line, end_line, "a page may skip only its boundary newline");
                }
            }
        }
        assert!(pages.iter().any(|page| page["partialLine"] == true), "the long first line must span actual pages");
        assert!(omitted_newline, "complete-line pagination must exercise its delimiter convention");
        assert_eq!(recovered.as_bytes(), expected.as_bytes(), "every original UTF-8 byte must be recoverable");
        assert_eq!(foreign, json!({"kind":"failed","message":"Tool result was not found in this Session"}));
        assert_eq!(restarted, recovery_results[0], "a new authorized Invocation after restart must read the same immutable page");
        let source_archives = original.events.iter().filter(|event| matches!(&event.event.fact,
            Fact::ToolResultArchived { placeholder } if placeholder.identity.runtime_event_id == source.event.id)).count();
        let archives = original.events.iter().filter(|event| matches!(&event.event.fact, Fact::ToolResultArchived { .. })).count();
        assert_eq!(archives, source_archives, "bounded recovery pages must not acquire their own archive replacements");
        assert_eq!(source_archives, 1);
        }
    }).await.unwrap();
}

#[tokio::test]
async fn todo_real_model_direct_and_code_mode_keep_full_list_isolation_and_restart() {
    tokio::time::timeout(Duration::from_secs(90), async {
        for code in [false, true] {
            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let mut f = fixture(&format!("http://{}/v1", listener.local_addr().unwrap())).await;
            let no_service_bindings = f.owner.tool_basis.is_empty();
            selection::edit(&f, |changes| {
                changes.model_overrides = Patch::Set(BTreeMap::from([("fixture-model".into(), ModelOverride { code_mode: Some(code), context_window: Some(1_000_000), input_limit: Some(900_000), ..Default::default() })]));
            }).await;
            let session = f.create("Checklist source").await;
            let other = f.create("Separate checklist").await;
            let items: Vec<Value> = (0..240).map(|index| json!({"content":format!("{index:03} {}", "😀".repeat(256)),"status":if index == 0 { "completed" } else if index == 1 { "in_progress" } else { "pending" }})).collect();
            f.send(&session).await.unwrap();
            let (mut socket, _) = listener.accept().await.unwrap();
            let sent = request(&mut socket).await;
            let tool = if code { "exec" } else { "todo_write" };
            assert!(sent["tools"].as_array().unwrap().iter().any(|t| t["function"]["name"] == tool));
            let input = if code { json!({"code":"const todos = Array.from({length:240}, (_, index) => ({content:String(index).padStart(3, '0') + ' ' + '😀'.repeat(256), status:index === 0 ? 'completed' : index === 1 ? 'in_progress' : 'pending'})); await tools.todo_write({todos}); const read = await tools.todo_read({}); if (JSON.stringify(read.items) !== JSON.stringify(todos)) throw Error('read differs from complete write'); text(read.items.at(-1));"}) } else { json!({"todos":items}) };
            service_tools::named_tool_response(&mut socket, tool, input).await;
            let (mut socket, _) = listener.accept().await.unwrap();
            let written = request(&mut socket).await;
            let direct_archive = if code {
                assert!(last_tool(&written).contains("239"), "Code Mode must read the complete list");
                None
            } else {
                // Ordinary large tool results use the existing bounded archive projection.
                let archived: Value = serde_json::from_str(last_tool(&written)).unwrap();
                assert_eq!(archived["kind"], "maka.archived_tool_result");
                assert_eq!(archived["toolName"], "todo_write");
                assert!(archived["page"]["next"].is_object());
                Some(archived)
            };
            if !code {
                service_tools::named_tool_response(&mut socket,"todo_read",json!({})).await;
                let (mut next, _) = listener.accept().await.unwrap();
                let read = request(&mut next).await;
                let archived: Value = serde_json::from_str(last_tool(&read)).unwrap();
                assert_eq!(archived["kind"], "maka.archived_tool_result");
                assert_eq!(archived["toolName"], "todo_read");
                response(&mut next).await;
            } else { response(&mut socket).await; }
            f.drained().await;
            let first = page(&f,&session,0).await;
            assert_eq!(first["total"],240);
            assert_eq!(first["completed"],1);
            assert_eq!(first["current"],items[1]);
            assert_eq!(first["items"],json!(items[..20]));
            assert!(serde_json::to_vec(&first).unwrap().len() < 48 * 1024);
            let mut all = vec![];
            for offset in (0..240).step_by(20) {
                let current = page(&f,&session,offset).await;
                all.extend(current["items"].as_array().unwrap().iter().cloned());
                assert_eq!(current["revision"],first["revision"]);
            }
            assert_eq!(all,items);
            assert_eq!(page(&f,&other,0).await["total"],0);
            assert!(f.query("checklist",json!({"sessionId":"foreign-session","offset":0})).await.is_err());
            let facts = service_tools::facts(&f,&session).await;
            let dispatched: Vec<_> = facts.iter().filter_map(|fact| match fact { Fact::ToolDispatched { name, operation_id, .. } if name == "todo_write" || name == "todo_read" => Some(operation_id), _ => None }).collect();
            assert_eq!(dispatched.len(),2);
            let original = f.store.scoped_prefix(maka_agent_facts::event::LogScope::Session { id: session.clone() }, 256, 4 * 1024 * 1024).await.unwrap();
            for id in dispatched {
                let event = original.events.iter().find(|event| matches!(&event.event.fact, Fact::ToolSettled { operation_id, .. } if operation_id == id)).unwrap();
                let output = f.store.resolve_tool_result(&session, &event.event.id).await.unwrap();
                let maka_agent_facts::tool_output::ToolOutput::Json(document) = output else { panic!("original checklist JSON is retained") };
                assert_eq!(document["items"],json!(items), "every item survives in each original write/read receipt");
                assert_eq!(facts.iter().filter(|fact| matches!(fact, Fact::ToolSettled { operation_id, outcome: ToolOutcome::Succeeded { .. } } if operation_id == id)).count(),1);
            }
            reopen(&mut f).await;
            assert_eq!(page(&f,&session,0).await,first);
            assert_eq!(page(&f,&other,0).await["total"],0);
            f.app.shutdown().await.unwrap();
            f.config.shutdown().await.unwrap();
            if let Some(archived) = direct_archive {
                // This fixture selected no service bindings, including Files under any alias.
                assert!(no_service_bindings);
                assert!(sent["tools"].as_array().unwrap().iter()
                    .all(|tool| tool["function"]["name"] != "Read"), "{sent}");
                assert!(archived["page"]["next"].is_object());
                assert_eq!(archived["readInstructions"],
                    "Use page.next only with a selected Session tool-result reader; without one, the remainder is unavailable.");
            }
        }
    }).await.unwrap();
}

#[tokio::test]
async fn todo_prepared_replacements_conflict_and_cancelled_or_closed_invocations_cannot_write() {
    tokio::time::timeout(Duration::from_secs(45), async {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let f = fixture(&format!("http://{}/v1", listener.local_addr().unwrap())).await;
        let session = f.create("Concurrent checklist").await;
        f.send(&session).await.unwrap();
        let (mut socket, _) = listener.accept().await.unwrap();
        request(&mut socket).await;
        let invocation = f
            .owner
            .runs
            .lock()
            .await
            .controls
            .get(&session)
            .unwrap()
            .invocation
            .clone();
        let journal = ToolJournal::new(f.store.clone(), invocation.clone());
        let context = |id: &str| ToolCallContext {
            invocation: invocation.clone(),
            operation_id: id.into(),
        };
        let input = |text: &str| json!({"todos":[{"content":text,"status":"in_progress"}]});
        let first = f
            .owner
            .todo
            .prepare(
                "todo_write".into(),
                input("First"),
                context("first"),
                CancellationToken::new(),
            )
            .await
            .unwrap();
        let second = f
            .owner
            .todo
            .prepare(
                "todo_write".into(),
                input("Second"),
                context("second"),
                CancellationToken::new(),
            )
            .await
            .unwrap();
        for (id, text, prepared, success) in [
            ("first", "First", first, true),
            ("second", "Second", second, false),
        ] {
            let outcome = journal
                .invoke_prepared_call(
                    id.into(),
                    ToolCallIdentity::standalone(id.into()),
                    "todo_write".into(),
                    input(text),
                    CancellationToken::new(),
                    prepared,
                )
                .await;
            assert_eq!(outcome.is_ok(), success, "{outcome:?}");
        }
        let before = page(&f, &session, 0).await;
        assert_eq!(before["items"][0]["content"], "First");
        let token = CancellationToken::new();
        token.cancel();
        assert!(matches!(
            f.owner
                .todo
                .prepare(
                    "todo_write".into(),
                    input("Cancelled"),
                    context("cancelled"),
                    token
                )
                .await,
            Err(ToolRejection::Cancelled)
        ));
        let pending = f
            .owner
            .todo
            .prepare(
                "todo_write".into(),
                input("Late"),
                context("late"),
                CancellationToken::new(),
            )
            .await
            .unwrap();
        response(&mut socket).await;
        f.drained().await;
        assert!(
            f.owner
                .todo
                .prepare(
                    "todo_write".into(),
                    input("Closed"),
                    context("closed"),
                    CancellationToken::new()
                )
                .await
                .is_err()
        );
        assert!(
            journal
                .invoke_prepared_call(
                    "late".into(),
                    ToolCallIdentity::standalone("late".into()),
                    "todo_write".into(),
                    input("Late"),
                    CancellationToken::new(),
                    pending
                )
                .await
                .is_err()
        );
        assert_eq!(page(&f, &session, 0).await, before);
        f.app.shutdown().await.unwrap();
    })
    .await
    .unwrap();
}