loong_kernel/
test_support.rs

1#![allow(clippy::expect_used, clippy::panic)]
2
3use std::collections::{BTreeMap, BTreeSet};
4use std::sync::Mutex;
5
6use async_trait::async_trait;
7use serde_json::json;
8
9use crate::connector::{ConnectorExtensionAdapter, CoreConnectorAdapter};
10use crate::contracts::{
11    Capability, ConnectorCommand, ConnectorOutcome, ExecutionRoute, HarnessKind, HarnessOutcome,
12    HarnessRequest,
13};
14use crate::errors::{ConnectorError, PolicyError};
15use crate::harness::HarnessAdapter;
16use crate::memory::{
17    CoreMemoryAdapter, MemoryCoreOutcome, MemoryCoreRequest, MemoryExtensionAdapter,
18    MemoryExtensionOutcome, MemoryExtensionRequest,
19};
20use crate::pack::VerticalPackManifest;
21use crate::policy_ext::{PolicyExtension, PolicyExtensionContext};
22use crate::runtime::{
23    CoreRuntimeAdapter, RuntimeCoreOutcome, RuntimeCoreRequest, RuntimeExtensionAdapter,
24    RuntimeExtensionOutcome, RuntimeExtensionRequest,
25};
26use crate::tool::{
27    CoreToolAdapter, ToolCoreOutcome, ToolCoreRequest, ToolExtensionAdapter, ToolExtensionOutcome,
28    ToolExtensionRequest,
29};
30
31pub struct MockEmbeddedPiHarness {
32    pub seen_tasks: Mutex<Vec<String>>,
33}
34pub struct MockCrmConnector;
35pub struct MockCoreConnector;
36pub struct MockCoreConnectorGrpc;
37pub struct MockPanickingCoreConnector;
38pub struct MockConnectorExtension;
39pub struct MockPanickingConnectorExtension;
40pub struct MockAcpHarness;
41pub struct MockCoreRuntime;
42pub struct MockCoreRuntimeFallback;
43pub struct MockRuntimeExtension;
44pub struct MockCoreTool;
45pub struct MockToolExtension;
46pub struct MockCoreMemory;
47pub struct MockMemoryExtension;
48pub struct NoNetworkEgressPolicyExtension;
49pub const TEST_CAPABILITY_VARIANTS: [Capability; 13] = [
50    Capability::InvokeTool,
51    Capability::InvokeConnector,
52    Capability::MemoryRead,
53    Capability::MemoryWrite,
54    Capability::FilesystemRead,
55    Capability::FilesystemWrite,
56    Capability::NetworkEgress,
57    Capability::ObserveTelemetry,
58    Capability::ControlRead,
59    Capability::ControlWrite,
60    Capability::ControlApprovals,
61    Capability::ControlPairing,
62    Capability::ControlAcp,
63];
64pub const TEST_CAPABILITY_VARIANT_COUNT: u8 = TEST_CAPABILITY_VARIANTS.len() as u8;
65#[derive(Debug, Clone, Copy)]
66pub enum ToolGateMode {
67    Deny,
68}
69#[derive(Debug)]
70pub struct ToolGatePolicyExtension {
71    pub gated_tool: String,
72    pub mode: ToolGateMode,
73}
74impl ToolGatePolicyExtension {
75    pub fn new(gated_tool: &str, mode: ToolGateMode) -> Self {
76        Self {
77            gated_tool: gated_tool.to_owned(),
78            mode,
79        }
80    }
81}
82
83#[async_trait]
84impl HarnessAdapter for MockEmbeddedPiHarness {
85    fn name(&self) -> &str {
86        "pi-local"
87    }
88    fn kind(&self) -> HarnessKind {
89        HarnessKind::EmbeddedPi
90    }
91    async fn execute(
92        &self,
93        request: HarnessRequest,
94    ) -> Result<HarnessOutcome, crate::HarnessError> {
95        self.seen_tasks
96            .lock()
97            .expect("mutex poisoned")
98            .push(request.task_id.clone());
99        Ok(HarnessOutcome {
100            status: "ok".to_owned(),
101            output: json!({"adapter":"pi-local","task_id":request.task_id,"objective":request.objective}),
102        })
103    }
104}
105#[async_trait]
106impl CoreConnectorAdapter for MockCrmConnector {
107    fn name(&self) -> &str {
108        "crm"
109    }
110    async fn invoke_core(
111        &self,
112        command: ConnectorCommand,
113    ) -> Result<ConnectorOutcome, ConnectorError> {
114        Ok(ConnectorOutcome {
115            status: "ok".to_owned(),
116            payload: json!({"operation":command.operation,"echo":command.payload}),
117        })
118    }
119}
120#[async_trait]
121impl CoreConnectorAdapter for MockCoreConnector {
122    fn name(&self) -> &str {
123        "http-core"
124    }
125    async fn invoke_core(
126        &self,
127        command: ConnectorCommand,
128    ) -> Result<ConnectorOutcome, ConnectorError> {
129        Ok(ConnectorOutcome {
130            status: "ok".to_owned(),
131            payload: json!({"tier":"core","adapter":"http-core","connector":command.connector_name,"operation":command.operation,"payload":command.payload}),
132        })
133    }
134}
135#[async_trait]
136impl CoreConnectorAdapter for MockCoreConnectorGrpc {
137    fn name(&self) -> &str {
138        "grpc-core"
139    }
140    async fn invoke_core(
141        &self,
142        command: ConnectorCommand,
143    ) -> Result<ConnectorOutcome, ConnectorError> {
144        Ok(ConnectorOutcome {
145            status: "ok".to_owned(),
146            payload: json!({"tier":"core","adapter":"grpc-core","connector":command.connector_name,"operation":command.operation}),
147        })
148    }
149}
150#[async_trait]
151impl CoreConnectorAdapter for MockPanickingCoreConnector {
152    fn name(&self) -> &str {
153        "panic-core"
154    }
155    async fn invoke_core(
156        &self,
157        _command: ConnectorCommand,
158    ) -> Result<ConnectorOutcome, ConnectorError> {
159        panic!("simulated connector core panic");
160    }
161}
162#[async_trait]
163impl ConnectorExtensionAdapter for MockConnectorExtension {
164    fn name(&self) -> &str {
165        "shielded-bridge"
166    }
167    async fn invoke_extension(
168        &self,
169        command: ConnectorCommand,
170        core: &(dyn CoreConnectorAdapter + Sync),
171    ) -> Result<ConnectorOutcome, ConnectorError> {
172        let core_probe = core
173            .invoke_core(ConnectorCommand {
174                connector_name: command.connector_name.clone(),
175                operation: "probe".to_owned(),
176                required_capabilities: BTreeSet::new(),
177                payload: json!({"mode":"probe"}),
178            })
179            .await?;
180        Ok(ConnectorOutcome {
181            status: "ok".to_owned(),
182            payload: json!({"tier":"extension","extension":"shielded-bridge","operation":command.operation,"core_probe":core_probe.payload,"payload":command.payload}),
183        })
184    }
185}
186#[async_trait]
187impl ConnectorExtensionAdapter for MockPanickingConnectorExtension {
188    fn name(&self) -> &str {
189        "panic-extension"
190    }
191    async fn invoke_extension(
192        &self,
193        _command: ConnectorCommand,
194        _core: &(dyn CoreConnectorAdapter + Sync),
195    ) -> Result<ConnectorOutcome, ConnectorError> {
196        panic!("simulated connector extension panic");
197    }
198}
199#[async_trait]
200impl HarnessAdapter for MockAcpHarness {
201    fn name(&self) -> &str {
202        "acp-gateway"
203    }
204    fn kind(&self) -> HarnessKind {
205        HarnessKind::Acp
206    }
207    async fn execute(
208        &self,
209        request: HarnessRequest,
210    ) -> Result<HarnessOutcome, crate::HarnessError> {
211        Ok(HarnessOutcome {
212            status: "ok".to_owned(),
213            output: json!({"adapter":"acp-gateway","task_id":request.task_id}),
214        })
215    }
216}
217#[async_trait]
218impl CoreRuntimeAdapter for MockCoreRuntime {
219    fn name(&self) -> &str {
220        "native-core"
221    }
222    async fn execute_core(
223        &self,
224        request: RuntimeCoreRequest,
225    ) -> Result<RuntimeCoreOutcome, crate::RuntimePlaneError> {
226        Ok(RuntimeCoreOutcome {
227            status: "ok".to_owned(),
228            payload: json!({"adapter":"native-core","action":request.action,"payload":request.payload}),
229        })
230    }
231}
232#[async_trait]
233impl CoreRuntimeAdapter for MockCoreRuntimeFallback {
234    fn name(&self) -> &str {
235        "fallback-core"
236    }
237    async fn execute_core(
238        &self,
239        request: RuntimeCoreRequest,
240    ) -> Result<RuntimeCoreOutcome, crate::RuntimePlaneError> {
241        Ok(RuntimeCoreOutcome {
242            status: "ok".to_owned(),
243            payload: json!({"adapter":"fallback-core","action":request.action}),
244        })
245    }
246}
247#[async_trait]
248impl RuntimeExtensionAdapter for MockRuntimeExtension {
249    fn name(&self) -> &str {
250        "acp-bridge"
251    }
252    async fn execute_extension(
253        &self,
254        request: RuntimeExtensionRequest,
255        core: &(dyn CoreRuntimeAdapter + Sync),
256    ) -> Result<RuntimeExtensionOutcome, crate::RuntimePlaneError> {
257        let core_probe = core
258            .execute_core(RuntimeCoreRequest {
259                action: "probe".to_owned(),
260                payload: json!({}),
261            })
262            .await?;
263        Ok(RuntimeExtensionOutcome {
264            status: "ok".to_owned(),
265            payload: json!({"extension":"acp-bridge","action":request.action,"core_probe":core_probe.payload,"payload":request.payload}),
266        })
267    }
268}
269#[async_trait]
270impl CoreToolAdapter for MockCoreTool {
271    fn name(&self) -> &str {
272        "core-tools"
273    }
274
275    async fn execute_core_tool(
276        &self,
277        request: ToolCoreRequest,
278    ) -> Result<ToolCoreOutcome, crate::ToolPlaneError> {
279        Ok(ToolCoreOutcome {
280            status: "ok".to_owned(),
281            payload: json!({"tool":request.tool_name,"payload":request.payload}),
282        })
283    }
284}
285#[async_trait]
286impl ToolExtensionAdapter for MockToolExtension {
287    fn name(&self) -> &str {
288        "sql-analytics"
289    }
290    async fn execute_tool_extension(
291        &self,
292        request: ToolExtensionRequest,
293        core: &(dyn CoreToolAdapter + Sync),
294    ) -> Result<ToolExtensionOutcome, crate::ToolPlaneError> {
295        let core_probe = core
296            .execute_core_tool(ToolCoreRequest {
297                tool_name: "schema_probe".to_owned(),
298                payload: json!({}),
299            })
300            .await?;
301        Ok(ToolExtensionOutcome {
302            status: "ok".to_owned(),
303            payload: json!({"extension":"sql-analytics","action":request.extension_action,"core_probe":core_probe.payload}),
304        })
305    }
306}
307#[async_trait]
308impl CoreMemoryAdapter for MockCoreMemory {
309    fn name(&self) -> &str {
310        "kv-core"
311    }
312    async fn execute_core_memory(
313        &self,
314        request: MemoryCoreRequest,
315    ) -> Result<MemoryCoreOutcome, crate::MemoryPlaneError> {
316        Ok(MemoryCoreOutcome {
317            status: "ok".to_owned(),
318            payload: json!({"operation":request.operation,"payload":request.payload}),
319        })
320    }
321}
322#[async_trait]
323impl MemoryExtensionAdapter for MockMemoryExtension {
324    fn name(&self) -> &str {
325        "vector-index"
326    }
327    async fn execute_memory_extension(
328        &self,
329        request: MemoryExtensionRequest,
330        core: &(dyn CoreMemoryAdapter + Sync),
331    ) -> Result<MemoryExtensionOutcome, crate::MemoryPlaneError> {
332        let core_probe = core
333            .execute_core_memory(MemoryCoreRequest {
334                operation: "read".to_owned(),
335                payload: json!({"key":"seed"}),
336            })
337            .await?;
338        Ok(MemoryExtensionOutcome {
339            status: "ok".to_owned(),
340            payload: json!({"extension":"vector-index","operation":request.operation,"core_probe":core_probe.payload}),
341        })
342    }
343}
344impl PolicyExtension for NoNetworkEgressPolicyExtension {
345    fn name(&self) -> &str {
346        "no-network-egress"
347    }
348    fn authorize_extension(&self, context: &PolicyExtensionContext<'_>) -> Result<(), PolicyError> {
349        if context
350            .required_capabilities
351            .contains(&Capability::NetworkEgress)
352        {
353            return Err(PolicyError::ExtensionDenied {
354                extension: self.name().to_owned(),
355                reason: "network egress is blocked for this environment".to_owned(),
356            });
357        }
358        Ok(())
359    }
360}
361impl PolicyExtension for ToolGatePolicyExtension {
362    fn name(&self) -> &str {
363        "tool-gate"
364    }
365    fn authorize_extension(&self, context: &PolicyExtensionContext<'_>) -> Result<(), PolicyError> {
366        let Some(params) = context.request_parameters else {
367            return Ok(());
368        };
369        let tool_name = params
370            .get("tool_name")
371            .and_then(|v| v.as_str())
372            .unwrap_or("");
373        if tool_name != self.gated_tool {
374            return Ok(());
375        }
376        match self.mode {
377            ToolGateMode::Deny => Err(PolicyError::ToolCallDenied {
378                tool_name: tool_name.to_owned(),
379                reason: "blocked by deterministic policy rule".to_owned(),
380            }),
381        }
382    }
383}
384pub fn sample_pack() -> VerticalPackManifest {
385    VerticalPackManifest {
386        pack_id: "sales-intel".to_owned(),
387        domain: "sales".to_owned(),
388        version: "0.1.0".to_owned(),
389        default_route: ExecutionRoute {
390            harness_kind: HarnessKind::EmbeddedPi,
391            adapter: Some("pi-local".to_owned()),
392        },
393        allowed_connectors: BTreeSet::from(["crm".to_owned()]),
394        granted_capabilities: BTreeSet::from([
395            Capability::InvokeTool,
396            Capability::InvokeConnector,
397            Capability::MemoryRead,
398        ]),
399        metadata: BTreeMap::from([("owner".to_owned(), "revenue-team".to_owned())]),
400    }
401}
402pub fn acp_pack_without_explicit_adapter() -> VerticalPackManifest {
403    VerticalPackManifest {
404        pack_id: "code-review".to_owned(),
405        domain: "engineering".to_owned(),
406        version: "0.1.0".to_owned(),
407        default_route: ExecutionRoute {
408            harness_kind: HarnessKind::Acp,
409            adapter: None,
410        },
411        allowed_connectors: BTreeSet::new(),
412        granted_capabilities: BTreeSet::from([Capability::InvokeTool]),
413        metadata: BTreeMap::new(),
414    }
415}
416pub fn capability_from_bit(bit: u8) -> Capability {
417    let bit_index = usize::from(bit);
418    TEST_CAPABILITY_VARIANTS
419        .get(bit_index)
420        .copied()
421        .expect("test capability bit should be in range")
422}
423pub fn capability_set_from_mask(mask: u16) -> BTreeSet<Capability> {
424    let mut capabilities = BTreeSet::new();
425    for (bit_index, capability) in TEST_CAPABILITY_VARIANTS.iter().copied().enumerate() {
426        let bit_mask = 1_u16 << bit_index;
427        let is_enabled = (mask & bit_mask) != 0;
428        if is_enabled {
429            capabilities.insert(capability);
430        }
431    }
432    capabilities
433}
loong_kernel/test_support.rs

loong_kernel/
test_support.rs
