johnandersen777.bsky.social/asyncai

Switch repository

Loading…

use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use async_trait::async_trait;
use asyncai::code_mode::{CodeModeLimits, CodeModeOptions};
use asyncai::{CodeModeToolExecutor, ToolCallOutput, ToolDefinition, ToolError, ToolExecutor};
use serde_json::{json, Value};

/// In-memory MCP stand-in that records every call.
#[derive(Default)]
struct FakeTools {
    calls: Mutex<Vec<(String, Value)>>,
}

impl FakeTools {
    fn calls(&self) -> Vec<(String, Value)> {
        self.calls.lock().unwrap().clone()
    }
}

#[async_trait]
impl ToolExecutor for FakeTools {
    async fn list_tools(&self) -> Result<Vec<ToolDefinition>, ToolError> {
        let object = |properties: Value, required: Value| json!({"type": "object", "properties": properties, "required": required});
        let none = None::<String>;
        let mut list_pods = ToolDefinition::new(
            "list_pods",
            Some("List pods in a namespace."),
            object(
                json!({"namespace": {"type": "string"}}),
                json!(["namespace"]),
            ),
        );
        list_pods.output_schema = Some(json!({
            "type": "object",
            "properties": {"items": {"type": "array", "items": {"type": "object", "properties": {
                "metadata": {"type": "object", "properties": {"name": {"type": "string"}}},
                "status": {"type": "object", "properties": {"phase": {"type": "string"}}},
            }}}},
        }));
        list_pods.read_only = Some(true);
        let mut delete_thing = ToolDefinition::new(
            "delete_thing",
            none.clone(),
            object(json!({"name": {"type": "string"}}), json!(["name"])),
        );
        delete_thing.destructive = Some(true);
        Ok(vec![
            list_pods,
            ToolDefinition::new(
                "get-events",
                Some("Events."),
                object(
                    json!({"namespace": {"type": "string"}}),
                    json!(["namespace"]),
                ),
            ),
            ToolDefinition::new("big", none.clone(), object(json!({}), json!([]))),
            ToolDefinition::new("fails", none.clone(), object(json!({}), json!([]))),
            ToolDefinition::new("logs", none.clone(), object(json!({}), json!([]))),
            ToolDefinition::new("slow", none, object(json!({}), json!([]))),
            delete_thing,
        ])
    }

    async fn call_tool(&self, name: &str, arguments: Value) -> Result<ToolCallOutput, ToolError> {
        self.calls
            .lock()
            .unwrap()
            .push((name.to_string(), arguments.clone()));
        let ok = |content: String| Ok(ToolCallOutput::text(content));
        match name {
            "list_pods" => {
                tokio::time::sleep(Duration::from_millis(200)).await;
                ok(json!({"items": [
                    {"metadata": {"name": "web-1"}, "status": {"phase": "Running"}},
                    {"metadata": {"name": "web-2"}, "status": {"phase": "Pending"}},
                    {"metadata": {"name": "db-1"}, "status": {"phase": "Failed"}},
                ]})
                .to_string())
            }
            "get-events" => {
                tokio::time::sleep(Duration::from_millis(200)).await;
                ok(json!({"items": [{"type": "Warning"}, {"type": "Normal"}, {"type": "Warning"}]})
                    .to_string())
            }
            "big" => {
                let items: Vec<Value> = (0..10_000)
                    .map(|index| json!({"name": format!("cm-{index}"), "data": "x".repeat(80)}))
                    .collect();
                ok(json!({"items": items}).to_string())
            }
            "fails" => Ok(ToolCallOutput::error("error: forbidden")),
            "logs" => ok("line one\nline two\n".to_string()),
            "slow" => {
                tokio::time::sleep(Duration::from_secs(30)).await;
                ok("too late".to_string())
            }
            "delete_thing" => ok(format!("deleted {}", arguments["name"])),
            other => Err(ToolError::new(format!("unknown tool {other}"))),
        }
    }
}

async fn code_mode_with(limits: CodeModeLimits) -> (Arc<FakeTools>, CodeModeToolExecutor) {
    let fake = Arc::new(FakeTools::default());
    let executor = CodeModeToolExecutor::with_limits(fake.clone(), limits, false)
        .await
        .unwrap();
    (fake, executor)
}

async fn code_mode() -> (Arc<FakeTools>, CodeModeToolExecutor) {
    code_mode_with(CodeModeLimits::default()).await
}

async fn run(executor: &dyn ToolExecutor, code: &str) -> ToolCallOutput {
    executor
        .call_tool("run_code", json!({"code": code}))
        .await
        .unwrap()
}

#[tokio::test]
async fn exposes_only_run_code_with_a_python_catalog() {
    let (_, executor) = code_mode().await;
    let tools = executor.list_tools().await.unwrap();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].name, "run_code");
    let description = tools[0].description.as_deref().unwrap();
    assert!(description.contains("async def list_pods(*, namespace: str) -> dict[str, Any]:"));
    assert!(description.contains(
        r#"Returns: {"items": [{"metadata": {"name": str, ...}, "status": {"phase": str, ...}, ...}], ...}"#
    ));
    assert!(description.contains("async def get_events(*, namespace: str) -> Any:"));
}

#[tokio::test]
async fn gathers_calls_concurrently_and_returns_only_the_filtered_value() {
    let (fake, executor) = code_mode().await;
    let started = Instant::now();
    let output = run(
        &executor,
        "import asyncio\n\
         pods, events = await asyncio.gather(list_pods(namespace='app'), get_events(namespace='app'))\n\
         bad = [p['metadata']['name'] for p in pods['items'] if p['status']['phase'] != 'Running']\n\
         {'not_running': bad, 'warnings': sum(1 for e in events['items'] if e['type'] == 'Warning')}",
    )
    .await;
    let elapsed = started.elapsed();

    assert!(!output.is_error, "{}", output.content);
    let value: Value = serde_json::from_str(&output.content).unwrap();
    assert_eq!(
        value,
        json!({"not_running": ["web-2", "db-1"], "warnings": 2})
    );

    let calls = fake.calls();
    assert_eq!(calls.len(), 2);
    assert!(calls.contains(&("list_pods".to_string(), json!({"namespace": "app"}))));
    assert!(calls.contains(&("get-events".to_string(), json!({"namespace": "app"}))));
    assert!(
        elapsed < Duration::from_millis(390),
        "gathered calls ran sequentially: {elapsed:?}"
    );
}

#[tokio::test]
async fn filters_a_megabyte_tool_result_inside_the_sandbox() {
    let (_, executor) = code_mode().await;
    let output = run(
        &executor,
        "data = await big()\n\
         cm_names = [item['name'] for item in data['items']]\n\
         {'count': len(cm_names), 'last': cm_names[-1]}",
    )
    .await;
    assert!(!output.is_error, "{}", output.content);
    assert_eq!(
        serde_json::from_str::<Value>(&output.content).unwrap(),
        json!({"count": 10000, "last": "cm-9999"})
    );
}

#[tokio::test]
async fn returning_too_much_is_rejected_but_the_variable_survives() {
    let (fake, executor) = code_mode().await;
    let run_scope = executor.run_scope().unwrap();

    let too_big = run(run_scope.as_ref(), "data = await big()\ndata").await;
    assert!(too_big.is_error);
    assert!(
        too_big.content.contains("character limit"),
        "{}",
        too_big.content
    );

    let filtered = run(run_scope.as_ref(), "len(data['items'])").await;
    assert!(!filtered.is_error, "{}", filtered.content);
    assert_eq!(filtered.content, "10000");
    assert_eq!(fake.calls().len(), 1, "variable reuse must not refetch");
}

#[tokio::test]
async fn tool_errors_raise_catchable_runtime_errors() {
    let (_, executor) = code_mode().await;
    let output = run(
        &executor,
        "try:\n    await fails()\n    result = 'no error'\nexcept RuntimeError as error:\n    result = 'caught: ' + str(error)\nresult",
    )
    .await;
    assert!(!output.is_error, "{}", output.content);
    assert_eq!(output.content, "caught: error: forbidden");

    let uncaught = run(&executor, "await fails()").await;
    assert!(uncaught.is_error);
    assert!(
        uncaught.content.contains("RuntimeError"),
        "{}",
        uncaught.content
    );
}

#[tokio::test]
async fn non_json_tool_output_arrives_as_str_and_prints_are_returned() {
    let (_, executor) = code_mode().await;
    let output = run(
        &executor,
        "text = await logs()\nprint('lines:', len(text.splitlines()))\ntext.splitlines()[0]",
    )
    .await;
    assert!(!output.is_error, "{}", output.content);
    assert_eq!(
        serde_json::from_str::<Value>(&output.content).unwrap(),
        json!({"output": "lines: 2\n", "result": "line one"})
    );
}

#[tokio::test]
async fn unknown_names_raise_name_error() {
    let (fake, executor) = code_mode().await;
    let output = run(&executor, "await kubectl_exec(cmd='rm -rf /')").await;
    assert!(output.is_error);
    assert!(output.content.contains("NameError"), "{}", output.content);
    assert!(fake.calls().is_empty());
}

#[tokio::test]
async fn os_access_is_denied() {
    let (_, executor) = code_mode().await;
    let read = run(
        &executor,
        "from pathlib import Path\nPath('/etc/passwd').read_text()",
    )
    .await;
    assert!(read.is_error);
    assert!(read.content.contains("PermissionError"), "{}", read.content);

    let env = run(&executor, "import os\nos.getenv('HOME')").await;
    assert!(env.is_error, "{}", env.content);
    assert!(env.content.contains("PermissionError"), "{}", env.content);
}

#[tokio::test]
async fn runaway_snippets_hit_the_time_limit() {
    let (_, executor) = code_mode_with(CodeModeLimits {
        max_duration: Duration::from_millis(200),
        ..CodeModeLimits::default()
    })
    .await;
    let started = Instant::now();
    let output = tokio::time::timeout(
        Duration::from_secs(10),
        run(&executor, "while True:\n    pass"),
    )
    .await
    .expect("run_code hung");
    assert!(output.is_error);
    assert!(output.content.contains("time limit"), "{}", output.content);
    assert!(started.elapsed() < Duration::from_secs(5));
}

#[tokio::test]
async fn repl_state_persists_within_a_scope_and_restart_clears_it() {
    let (_, executor) = code_mode().await;
    let scope = executor.run_scope().unwrap();

    assert_eq!(
        run(scope.as_ref(), "counter = 41").await.content,
        asyncai::code_mode::EMPTY_RESULT
    );
    assert_eq!(run(scope.as_ref(), "counter + 1").await.content, "42");

    let restarted = scope
        .call_tool("run_code", json!({"code": "counter", "restart": true}))
        .await
        .unwrap();
    assert!(restarted.is_error);
    assert!(
        restarted.content.contains("NameError"),
        "{}",
        restarted.content
    );

    let other_scope = executor.run_scope().unwrap();
    run(scope.as_ref(), "shared = 1").await;
    let isolated = run(other_scope.as_ref(), "shared").await;
    assert!(
        isolated.is_error,
        "scopes leaked state: {}",
        isolated.content
    );
}

#[tokio::test]
async fn positional_arguments_raise_type_error() {
    let (fake, executor) = code_mode().await;
    let output = run(&executor, "await list_pods('app')").await;
    assert!(output.is_error);
    assert!(output.content.contains("TypeError"), "{}", output.content);
    assert!(output.content.contains("keyword"), "{}", output.content);
    assert!(fake.calls().is_empty());
}

#[tokio::test]
async fn nested_call_budget_is_enforced() {
    let (fake, executor) = code_mode_with(CodeModeLimits {
        max_tool_calls: 3,
        ..CodeModeLimits::default()
    })
    .await;
    let output = run(
        &executor,
        "done = 0\nfor _ in range(5):\n    await logs()\n    done += 1\ndone",
    )
    .await;
    assert!(output.is_error);
    assert!(output.content.contains("budget"), "{}", output.content);
    assert_eq!(fake.calls().len(), 3);
}

#[tokio::test]
async fn unknown_and_missing_keywords_raise_type_error_before_dispatch() {
    let (fake, executor) = code_mode().await;
    let unknown = run(&executor, "await list_pods(namespace='a', name='x')").await;
    assert!(unknown.is_error);
    assert!(
        unknown
            .content
            .contains("unexpected keyword argument 'name'; accepted: namespace"),
        "{}",
        unknown.content
    );
    let missing = run(&executor, "await list_pods()").await;
    assert!(missing
        .content
        .contains("missing required keyword argument 'namespace'"));
    let none_optional = run(&executor, "await delete_thing(name='a')").await;
    assert!(!none_optional.is_error, "{}", none_optional.content);
    assert_eq!(fake.calls().len(), 1, "invalid calls must not reach MCP");
}

#[tokio::test]
async fn slow_tools_time_out_inside_the_snippet() {
    let (_, executor) = code_mode_with(CodeModeLimits {
        tool_timeout: Duration::from_millis(200),
        ..CodeModeLimits::default()
    })
    .await;
    let started = Instant::now();
    let output = run(
        &executor,
        "try:\n    await slow()\n    r = 'finished'\nexcept RuntimeError as e:\n    r = str(e)\nr",
    )
    .await;
    assert!(!output.is_error, "{}", output.content);
    assert!(
        output.content.contains("timed out after 200ms"),
        "{}",
        output.content
    );
    assert!(started.elapsed() < Duration::from_secs(5));
}

#[tokio::test]
async fn wall_deadline_ends_run_code() {
    let (_, executor) = code_mode_with(CodeModeLimits {
        max_wall: Duration::from_millis(300),
        ..CodeModeLimits::default()
    })
    .await;
    let started = Instant::now();
    let output = run(&executor, "await slow()").await;
    assert!(output.is_error);
    assert!(
        output.content.contains("wall-clock limit"),
        "{}",
        output.content
    );
    assert!(started.elapsed() < Duration::from_millis(600));
    let audit = output.audit.expect("audit recorded");
    assert_eq!(audit["nested_calls"][0]["tool"], "slow");
    assert_eq!(audit["nested_calls"][0]["finished"], false);
}

#[tokio::test]
async fn audit_records_nested_calls_and_mutating_flag() {
    let (_, executor) = code_mode().await;
    let output = run(
        &executor,
        "await delete_thing(name='cm-1')\nlen(items(await list_pods(namespace='a')))",
    )
    .await;
    assert!(!output.is_error, "{}", output.content);
    assert_eq!(output.content, "3");
    let calls = output.audit.unwrap()["nested_calls"].clone();
    assert_eq!(calls[0]["tool"], "delete_thing");
    assert_eq!(calls[0]["mutating"], true);
    assert_eq!(calls[0]["arguments"], json!({"name": "cm-1"}));
    assert_eq!(calls[1]["tool"], "list_pods");
    assert_eq!(calls[1]["mutating"], false);
    assert!(calls[1]["bytes"].as_u64().unwrap() > 100);
}

#[tokio::test]
async fn builtin_helpers_and_preamble_are_available() {
    let fake = Arc::new(FakeTools::default());
    let executor = CodeModeToolExecutor::with_options(
        fake,
        CodeModeOptions {
            preamble: Some("NAMESPACE = 'app'".to_string()),
            ..CodeModeOptions::default()
        },
    )
    .await
    .unwrap();
    let scope = executor.run_scope().unwrap();
    let output = run(
        scope.as_ref(),
        "pods, _ = await asyncio.gather(list_pods(namespace=NAMESPACE), list_pods(namespace=NAMESPACE))\n\
         {'names': names(pods), 'phases': count_by(pods, 'status.phase'), 'first': dig(pods, 'items.0.metadata.name')}",
    )
    .await;
    assert!(!output.is_error, "{}", output.content);
    assert_eq!(
        serde_json::from_str::<Value>(&output.content).unwrap(),
        json!({
            "names": ["web-1", "web-2", "db-1"],
            "phases": {"Running": 1, "Pending": 1, "Failed": 1},
            "first": "web-1",
        })
    );

    let restarted = scope
        .call_tool(
            "run_code",
            json!({"code": "[NAMESPACE, items(None)]", "restart": true}),
        )
        .await
        .unwrap();
    assert_eq!(
        restarted.content, r#"["app",[]]"#,
        "helpers re-fed after restart"
    );
}

#[tokio::test]
async fn non_run_code_tool_names_are_rejected() {
    let (fake, executor) = code_mode().await;
    let output = executor
        .call_tool("list_pods", json!({"namespace": "app"}))
        .await
        .unwrap();
    assert!(output.is_error);
    assert!(output.content.contains("run_code"));
    assert!(fake.calls().is_empty());
}

#[tokio::test]
async fn thread_scopes_persist_across_runs_and_isolate_threads() {
    let (_, executor) = code_mode().await;
    let first_run = executor.thread_scope("thread-1").unwrap();
    run(first_run.as_ref(), "seen = ['default/a']").await;
    drop(first_run);

    let next_run = executor.thread_scope("thread-1").unwrap();
    assert_eq!(
        run(next_run.as_ref(), "seen").await.content,
        r#"["default/a"]"#
    );

    let other_thread = executor.thread_scope("thread-2").unwrap();
    let isolated = run(other_thread.as_ref(), "seen").await;
    assert!(
        isolated.is_error,
        "threads leaked state: {}",
        isolated.content
    );
}

#[tokio::test]
async fn dig_takes_split_paths_and_to_json_stringifies() {
    let (_, executor) = code_mode().await;
    let scope = executor.run_scope().unwrap();
    let output = run(
        scope.as_ref(),
        "o = {'spec': {'payload': {'requesterDid': 'did:a'}}, 'items': [{'status': {'phase': 'Up'}}]}\n\
         [dig(o, 'spec.payload', 'requesterDid'), dig(o, 'spec', 'payload', 'requesterDid'), \
          dig(o, 'spec.missing', default='d'), dig(o, 'spec.missing', {}), \
          count_by(o, 'status', 'phase'), to_json({'k': {1, 2}} if False else {'k': (1, 2)})]",
    )
    .await;
    assert!(!output.is_error, "{}", output.content);
    assert_eq!(
        serde_json::from_str::<Value>(&output.content).unwrap(),
        json!(["did:a", "did:a", "d", {}, {"Up": 1}, "{\"k\": [1, 2]}"])
    );
}

#[tokio::test]
async fn helper_names_are_reserved() {
    let (_, executor) = code_mode().await;
    let scope = executor.run_scope().unwrap();
    let rejected = run(scope.as_ref(), "names = []\nnames").await;
    assert!(rejected.is_error);
    assert!(rejected.content.contains("`names` is a built-in helper"));
    let still_helper = run(scope.as_ref(), "names([{'name': 'a'}])").await;
    assert_eq!(still_helper.content, r#"["a"]"#);
}

#[tokio::test]
async fn unawaited_gather_and_empty_results_are_explicit() {
    let (_, executor) = code_mode().await;
    let scope = executor.run_scope().unwrap();
    let unawaited = run(
        scope.as_ref(),
        "asyncio.gather(list_pods(namespace='a'), list_pods(namespace='b'))",
    )
    .await;
    assert!(unawaited.is_error, "{}", unawaited.content);
    assert!(unawaited.content.contains("add `await`"));

    let empty = run(scope.as_ref(), "x = 1").await;
    assert!(!empty.is_error);
    assert_eq!(empty.content, asyncai::code_mode::EMPTY_RESULT);

    let hinted = run(scope.as_ref(), "json.dumps({'a': 1}, default=str)").await;
    assert!(hinted.is_error);
    assert!(hinted.content.contains("Hint:"), "{}", hinted.content);
}

#[tokio::test]
async fn thread_scope_notes_list_variables_still_bound() {
    let (_, executor) = code_mode().await;
    let scope = executor.thread_scope("t").unwrap();
    assert_eq!(scope.context_note(), None);
    run(
        scope.as_ref(),
        "pods = await list_pods(namespace='a')\nids = [1, 2, 3]",
    )
    .await;
    run(scope.as_ref(), "label = 'web'").await;
    let note = executor.thread_scope("t").unwrap().context_note().unwrap();
    assert!(note.contains("pods: dict["), "{note}");
    assert!(note.contains("label: str[3]"), "{note}");
    assert!(note.contains("ids: list[3]"), "{note}");

    run(scope.as_ref(), "x = 1").await;
    scope
        .call_tool("run_code", json!({"code": "y = 2", "restart": true}))
        .await
        .unwrap();
    let note = scope.context_note().unwrap();
    assert!(note.contains("y: int") && !note.contains("pods"), "{note}");
}