starweaver-runtime 0.0.2

Agent-loop graph and runtime executor primitives for Starweaver
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

//! Dependency-aware capability hook tests.

use std::sync::{Arc, Mutex};

use async_trait::async_trait;
use starweaver_context::{AgentContext, BusMessage};
use starweaver_model::{ModelRequest, ModelResponse, ModelSettings, TestModel};
use starweaver_runtime::{Agent, AgentCapability, AgentRunState, CapabilityResult};

#[derive(Clone, Debug)]
struct Tenant(String);

struct ContextAwareCapability {
    observed: Arc<Mutex<Vec<String>>>,
}

#[async_trait]
impl AgentCapability for ContextAwareCapability {
    async fn on_run_start_with_context(
        &self,
        _state: &mut AgentRunState,
        context: &mut AgentContext,
    ) -> CapabilityResult<()> {
        let tenant = context
            .dependency::<Tenant>()
            .map_or_else(|| "missing".to_string(), |tenant| tenant.0.clone());
        if let Ok(mut observed) = self.observed.lock() {
            observed.push(format!("start:{tenant}"));
        }
        context.enqueue_message(BusMessage::new(
            "capability.started",
            serde_json::json!({"tenant": tenant}),
        ));
        Ok(())
    }

    async fn before_model_request_with_context(
        &self,
        _state: &mut AgentRunState,
        context: &mut AgentContext,
        request: &mut ModelRequest,
        _settings: &mut Option<ModelSettings>,
    ) -> CapabilityResult<()> {
        let tenant = context
            .dependency::<Tenant>()
            .map_or_else(|| "missing".to_string(), |tenant| tenant.0.clone());
        request
            .metadata
            .insert("tenant".to_string(), serde_json::json!(tenant));
        Ok(())
    }

    async fn after_model_response_with_context(
        &self,
        _state: &mut AgentRunState,
        context: &mut AgentContext,
        response: &mut ModelResponse,
    ) -> CapabilityResult<()> {
        let tenant = context
            .dependency::<Tenant>()
            .map_or_else(|| "missing".to_string(), |tenant| tenant.0.clone());
        response
            .metadata
            .insert("tenant".to_string(), serde_json::json!(tenant));
        Ok(())
    }

    async fn on_run_complete_with_context(
        &self,
        _state: &mut AgentRunState,
        context: &mut AgentContext,
    ) -> CapabilityResult<()> {
        if let Ok(mut observed) = self.observed.lock() {
            observed.push(format!("messages:{}", context.messages.len()));
        }
        Ok(())
    }
}

#[tokio::test]
async fn capability_hooks_can_read_context_dependencies() {
    let observed = Arc::new(Mutex::new(Vec::new()));
    let capability = Arc::new(ContextAwareCapability {
        observed: observed.clone(),
    });
    let model = Arc::new(TestModel::with_text("ok"));
    let mut context = AgentContext::default();
    context.insert_dependency(Tenant("acme".to_string()));

    let result = Agent::new(model)
        .with_capability(capability)
        .run_with_context("hello", &mut context)
        .await;

    assert!(result.is_ok());
    assert_eq!(
        observed
            .lock()
            .map_or_else(|_| Vec::new(), |observed| observed.clone()),
        vec!["start:acme".to_string(), "messages:1".to_string()]
    );
    assert_eq!(context.messages.len(), 1);
}