a3s-code-core 8.5.5

A3S Code Core - Embeddable AI agent library with tool execution
Documentation
use super::*;
use crate::evaluation::evidence::{EvidenceReadRequestV1, RunEvidenceReader};
use crate::evaluation::identity::{digest_bytes, ExecutionTargetV1};
use crate::run::InMemoryRunStore;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::Duration;

struct FixtureExecutor {
    calls: AtomicUsize,
    value: serde_json::Value,
}

#[async_trait]
impl AuxiliaryExecutor for FixtureExecutor {
    async fn execute(
        &self,
        context: AuxiliaryRunContextV1,
    ) -> Result<serde_json::Value, AuxiliaryRunError> {
        self.calls.fetch_add(1, Ordering::SeqCst);
        assert!(!context.spec.capabilities.write_workspace);
        Ok(self.value.clone())
    }
}

async fn evidence() -> EvidenceSnapshotV1 {
    let runs = Arc::new(InMemoryRunStore::new());
    let run = runs
        .create_run_with_id("run-1".into(), "session-1", "prompt")
        .await;
    RunEvidenceReader::new(runs)
        .read(EvidenceReadRequestV1::new(ExecutionTargetV1::new(
            "session-1",
            &run.id,
        )))
        .await
        .unwrap()
}

#[tokio::test]
async fn service_runs_isolated_executor_and_is_idempotent() {
    let evidence = evidence().await;
    let executor = Arc::new(FixtureExecutor {
        calls: AtomicUsize::new(0),
        value: serde_json::json!({"decision": "ok"}),
    });
    let service = InMemoryAuxiliaryRunService::new(executor.clone());
    let spec = AuxiliaryRunSpecV1::new(
        ExecutionFrameV1::root(evidence.target.clone()),
        "fixture",
        "inspect bounded evidence",
        evidence.snapshot_digest.clone(),
    )
    .with_id("aux-1")
    .with_capabilities(AuxiliaryCapabilityProfileV1::tool_free());
    let handle = service
        .spawn(spec.clone(), evidence.clone(), None)
        .await
        .unwrap();
    let replay = service.spawn(spec, evidence, None).await.unwrap();
    assert_eq!(handle.id(), replay.id());
    let output = handle.wait().await.unwrap();
    assert_eq!(output.value["decision"], "ok");
    assert_eq!(executor.calls.load(Ordering::SeqCst), 1);
    assert_eq!(
        service.get("aux-1").await.unwrap().state,
        AuxiliaryRunStateV1::Completed
    );
}

#[tokio::test]
async fn parent_cancellation_is_propagated() {
    struct BlockingExecutor;
    #[async_trait]
    impl AuxiliaryExecutor for BlockingExecutor {
        async fn execute(
            &self,
            context: AuxiliaryRunContextV1,
        ) -> Result<serde_json::Value, AuxiliaryRunError> {
            context.cancellation.cancelled().await;
            Err(AuxiliaryRunError::Cancelled)
        }
    }
    let evidence = evidence().await;
    let service = InMemoryAuxiliaryRunService::new(Arc::new(BlockingExecutor));
    let parent = CancellationToken::new();
    let spec = AuxiliaryRunSpecV1::new(
        ExecutionFrameV1::root(evidence.target.clone()),
        "cancel-fixture",
        "wait",
        evidence.snapshot_digest.clone(),
    )
    .with_id("aux-cancel");
    let handle = service
        .spawn(spec, evidence, Some(parent.clone()))
        .await
        .unwrap();
    parent.cancel();
    assert!(matches!(
        handle.wait().await,
        Err(AuxiliaryRunError::Cancelled)
    ));
    assert_eq!(
        handle.snapshot().await.state,
        AuxiliaryRunStateV1::Cancelled
    );
}

#[test]
fn capability_ceiling_rejects_escalation() {
    let target = ExecutionTargetV1::new("s", "r");
    let frame = ExecutionFrameV1::root(target);
    let digest = digest_bytes("evidence", b"x");
    let spec = AuxiliaryRunSpecV1::new(frame, "x", "y", digest)
        .with_capabilities(AuxiliaryCapabilityProfileV1::read_only(1024))
        .with_parent_ceiling(AuxiliaryCapabilityProfileV1::tool_free());
    assert!(matches!(
        spec.validate(&spec.evidence_digest),
        Err(AuxiliaryRunError::CapabilityEscalation)
    ));
}

#[tokio::test]
async fn output_schema_and_limit_are_enforced() {
    let evidence = evidence().await;
    let service = InMemoryAuxiliaryRunService::new(Arc::new(FixtureExecutor {
        calls: AtomicUsize::new(0),
        value: serde_json::json!({"wrong": true}),
    }));
    let spec = AuxiliaryRunSpecV1::new(
        ExecutionFrameV1::root(evidence.target.clone()),
        "schema-fixture",
        "return object",
        evidence.snapshot_digest.clone(),
    )
    .with_id("aux-schema")
    .with_output_schema(serde_json::json!({
        "type": "object",
        "required": ["answer"],
        "properties": {"answer": {"type": "string"}}
    }));
    let handle = service.spawn(spec, evidence, None).await.unwrap();
    assert!(matches!(
        handle.wait().await,
        Err(AuxiliaryRunError::OutputSchemaMismatch)
    ));
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_waiters_cannot_miss_terminal_notification() {
    struct YieldingExecutor;
    #[async_trait]
    impl AuxiliaryExecutor for YieldingExecutor {
        async fn execute(
            &self,
            _context: AuxiliaryRunContextV1,
        ) -> Result<serde_json::Value, AuxiliaryRunError> {
            // Give waiter tasks a chance to enter the check/register window.
            tokio::task::yield_now().await;
            Ok(serde_json::json!({"ok": true}))
        }
    }

    let evidence = evidence().await;
    let service = InMemoryAuxiliaryRunService::new(Arc::new(YieldingExecutor));
    for index in 0..256 {
        let spec = AuxiliaryRunSpecV1::new(
            ExecutionFrameV1::root(evidence.target.clone()),
            "wait-race",
            "return",
            evidence.snapshot_digest.clone(),
        )
        .with_id(format!("aux-wait-race-{index}"));
        let handle = service.spawn(spec, evidence.clone(), None).await.unwrap();
        let waiters = (0..4)
            .map(|_| {
                let handle = handle.clone();
                tokio::spawn(async move { handle.wait().await })
            })
            .collect::<Vec<_>>();
        for waiter in waiters {
            let result = tokio::time::timeout(Duration::from_secs(1), waiter)
                .await
                .expect("terminal notification must not be lost")
                .expect("waiter task must not panic")
                .expect("fixture executor succeeds");
            assert_eq!(result.value["ok"], true);
        }
        assert!(handle.snapshot().await.state.is_terminal());
    }
}

#[test]
fn auxiliary_snapshot_state_and_terminal_fields_are_consistent() {
    let target = ExecutionTargetV1::new("session-snapshot", "run-snapshot");
    let frame = ExecutionFrameV1::root(target);
    let spec = AuxiliaryRunSpecV1::new(
        frame.clone(),
        "snapshot",
        "return",
        digest_bytes("evidence", b"snapshot"),
    );
    let mut snapshot = AuxiliaryRunSnapshotV1 {
        schema: AUXILIARY_SNAPSHOT_SCHEMA_V1.to_string(),
        id: spec.id.clone(),
        parent: frame,
        mode: spec.mode,
        state: AuxiliaryRunStateV1::Failed,
        spec_digest: spec.digest().unwrap(),
        created_at_ms: 1,
        updated_at_ms: 1,
        output_digest: None,
        error: None,
    };
    assert!(matches!(
        snapshot.validate(),
        Err(AuxiliaryRunError::InvalidField("state"))
    ));
    snapshot.error = Some("failed".into());
    assert!(snapshot.validate().is_ok());
}

#[tokio::test]
async fn service_rejects_cross_target_evidence() {
    let evidence = evidence().await;
    let other_target = ExecutionTargetV1::new("session-other", "run-other");
    let spec = AuxiliaryRunSpecV1::new(
        ExecutionFrameV1::root(other_target),
        "cross-target",
        "must be rejected",
        evidence.snapshot_digest.clone(),
    )
    .with_id("aux-cross-target");
    let service = InMemoryAuxiliaryRunService::new(Arc::new(FixtureExecutor {
        calls: AtomicUsize::new(0),
        value: serde_json::json!({}),
    }));
    assert!(matches!(
        service.spawn(spec, evidence, None).await,
        Err(AuxiliaryRunError::TargetMismatch)
    ));
}

#[tokio::test]
async fn gate_mode_rejects_incomplete_evidence_while_advisory_may_proceed() {
    use crate::agent::AgentEvent;
    use crate::evaluation::evidence::EvidenceContentModeV1;

    let runs = Arc::new(InMemoryRunStore::new());
    let run = runs
        .create_run_with_id("run-incomplete".into(), "session-1", "prompt")
        .await;
    runs.record_event(
        &run.id,
        AgentEvent::TextDelta {
            text: "oversized evidence payload".into(),
        },
    )
    .await;
    let mut request = EvidenceReadRequestV1::new(ExecutionTargetV1::new("session-1", &run.id));
    request.content_mode = EvidenceContentModeV1::BoundedPayload;
    request.limits.max_event_bytes = 1;
    let incomplete = RunEvidenceReader::new(runs).read(request).await.unwrap();
    incomplete.validate().unwrap();
    assert!(!incomplete.complete);

    let service = InMemoryAuxiliaryRunService::new(Arc::new(FixtureExecutor {
        calls: AtomicUsize::new(0),
        value: serde_json::json!({"decision": "ok"}),
    }));
    let gate = AuxiliaryRunSpecV1::new(
        ExecutionFrameV1::root(incomplete.target.clone()),
        "gate-incomplete",
        "must fail closed",
        incomplete.snapshot_digest.clone(),
    )
    .with_id("aux-gate-incomplete")
    .with_mode(AuxiliaryModeV1::Gate);
    assert!(matches!(
        service.spawn(gate, incomplete.clone(), None).await,
        Err(AuxiliaryRunError::EvidenceIncomplete)
    ));

    let advisory = AuxiliaryRunSpecV1::new(
        ExecutionFrameV1::root(incomplete.target.clone()),
        "advisory-incomplete",
        "host may inspect incomplete evidence",
        incomplete.snapshot_digest.clone(),
    )
    .with_id("aux-advisory-incomplete")
    .with_mode(AuxiliaryModeV1::Advisory);
    let handle = service.spawn(advisory, incomplete, None).await.unwrap();
    assert_eq!(
        handle.wait().await.unwrap().value["decision"],
        serde_json::json!("ok")
    );
}

#[tokio::test]
async fn executor_panics_and_output_overflow_become_terminal_failures() {
    struct PanicExecutor;
    #[async_trait]
    impl AuxiliaryExecutor for PanicExecutor {
        async fn execute(
            &self,
            _context: AuxiliaryRunContextV1,
        ) -> Result<serde_json::Value, AuxiliaryRunError> {
            panic!("fixture panic");
        }
    }

    let evidence = evidence().await;
    let service = InMemoryAuxiliaryRunService::new(Arc::new(PanicExecutor));
    let spec = AuxiliaryRunSpecV1::new(
        ExecutionFrameV1::root(evidence.target.clone()),
        "panic",
        "panic",
        evidence.snapshot_digest.clone(),
    )
    .with_id("aux-panic");
    let handle = service.spawn(spec, evidence.clone(), None).await.unwrap();
    assert!(matches!(
        handle.wait().await,
        Err(AuxiliaryRunError::Executor(message)) if message.contains("panicked")
    ));
    assert_eq!(handle.snapshot().await.state, AuxiliaryRunStateV1::Failed);

    struct LargeExecutor;
    #[async_trait]
    impl AuxiliaryExecutor for LargeExecutor {
        async fn execute(
            &self,
            _context: AuxiliaryRunContextV1,
        ) -> Result<serde_json::Value, AuxiliaryRunError> {
            Ok(serde_json::json!("too large"))
        }
    }
    let service = InMemoryAuxiliaryRunService::new(Arc::new(LargeExecutor));
    let spec = AuxiliaryRunSpecV1::new(
        ExecutionFrameV1::root(evidence.target.clone()),
        "overflow",
        "overflow",
        evidence.snapshot_digest.clone(),
    )
    .with_id("aux-overflow")
    .with_capabilities(AuxiliaryCapabilityProfileV1::read_only(1));
    let handle = service.spawn(spec, evidence, None).await.unwrap();
    assert!(matches!(
        handle.wait().await,
        Err(AuxiliaryRunError::OutputLimit)
    ));
    assert_eq!(handle.snapshot().await.state, AuxiliaryRunStateV1::Failed);
}