agentsight_ext_runtime/
lib.rs1use wasmtime::component::{Component, Linker, ResourceTable};
5use wasmtime::{Config, Engine, Store, StoreLimits, StoreLimitsBuilder};
6use wasmtime_wasi::{WasiCtx, WasiCtxView, WasiView};
7
8const DEFAULT_MEMORY_BYTES: usize = 64 * 1024 * 1024;
9const DEFAULT_FUEL: u64 = 10_000_000;
10const MAX_CORE_INSTANCES: usize = 16;
11const MAX_MEMORIES: usize = 4;
12const MAX_TABLES: usize = 8;
13const MAX_COMPONENT_BYTES: usize = 16 * 1024 * 1024;
14const MAX_CONTENT_BYTES: usize = 16 * 1024 * 1024;
15const MAX_METADATA_BYTES: usize = 64 * 1024;
16
17struct ExtStore {
18 limits: StoreLimits,
19 table: ResourceTable,
20 wasi: WasiCtx,
21}
22
23impl WasiView for ExtStore {
24 fn ctx(&mut self) -> WasiCtxView<'_> {
25 WasiCtxView {
26 ctx: &mut self.wasi,
27 table: &mut self.table,
28 }
29 }
30}
31
32pub struct ExtRuntime {
43 engine: Engine,
44}
45
46impl ExtRuntime {
47 pub fn new() -> Result<Self, wasmtime::Error> {
48 let mut config = Config::new();
49 config.consume_fuel(true);
50 Ok(Self {
51 engine: Engine::new(&config)?,
52 })
53 }
54
55 fn store(&self) -> Result<Store<ExtStore>, wasmtime::Error> {
56 let mut wasi = WasiCtx::builder();
57 wasi.allow_tcp(false).allow_udp(false);
58 let mut store = Store::new(
59 &self.engine,
60 ExtStore {
61 limits: StoreLimitsBuilder::new()
62 .memory_size(DEFAULT_MEMORY_BYTES)
63 .instances(MAX_CORE_INSTANCES)
64 .memories(MAX_MEMORIES)
65 .tables(MAX_TABLES)
66 .build(),
67 table: ResourceTable::new(),
68 wasi: wasi.build(),
69 },
70 );
71 store.limiter(|state| &mut state.limits);
72 store.set_fuel(DEFAULT_FUEL)?;
73 Ok(store)
74 }
75
76 pub fn session_parse(
77 &self,
78 component_bytes: &[u8],
79 agent: &str,
80 path: &str,
81 updated_ms: u64,
82 content: &str,
83 ) -> Result<Option<String>, wasmtime::Error> {
84 validate_session_input(component_bytes, agent, path, content)?;
85 let component = Component::from_binary(&self.engine, component_bytes)?;
86 let mut linker = Linker::<ExtStore>::new(&self.engine);
87 wasmtime_wasi::p2::add_to_linker_sync(&mut linker)?;
88 let mut store = self.store()?;
89 let instance = linker.instantiate(&mut store, &component)?;
90 let parse = instance.get_typed_func::<(String, String, u64, String), (Option<String>,)>(
91 &mut store, "parse",
92 )?;
93 Ok(parse
94 .call(
95 &mut store,
96 (
97 agent.to_owned(),
98 path.to_owned(),
99 updated_ms,
100 content.to_owned(),
101 ),
102 )?
103 .0)
104 }
105}
106
107fn validate_session_input(
108 component_bytes: &[u8],
109 agent: &str,
110 path: &str,
111 content: &str,
112) -> Result<(), wasmtime::Error> {
113 if component_bytes.len() > MAX_COMPONENT_BYTES {
114 return Err(wasmtime::Error::msg("extension component exceeds 16 MiB"));
115 }
116 if content.len() > MAX_CONTENT_BYTES {
117 return Err(wasmtime::Error::msg("session content exceeds 16 MiB"));
118 }
119 if agent.len().saturating_add(path.len()) > MAX_METADATA_BYTES {
120 return Err(wasmtime::Error::msg("session metadata exceeds 64 KiB"));
121 }
122 Ok(())
123}
124
125#[cfg(test)]
126mod tests {
127 use super::*;
128
129 #[test]
130 fn rejects_oversized_component_content_and_metadata() {
131 assert!(
132 validate_session_input(&vec![0; MAX_COMPONENT_BYTES + 1], "codex", "x", "x").is_err()
133 );
134 assert!(
135 validate_session_input(&[], "codex", "x", &"x".repeat(MAX_CONTENT_BYTES + 1)).is_err()
136 );
137 assert!(validate_session_input(&[], &"a".repeat(MAX_METADATA_BYTES), "x", "x").is_err());
138 }
139}