feat: per-user auto-memory with ResponseContext
Three memory channels: hidden tool (sol.memory.set/get in scripts), pre-response injection (relevant memories loaded into system prompt), and post-response extraction (ministral-3b extracts facts after each response). User isolation enforced at Rust level — user_id derived from Matrix sender, never from script arguments. New modules: context (ResponseContext), memory (schema, store, extractor). ResponseContext threaded through responder → tools → script runtime. OpenSearch index sol_user_memory created on startup alongside archive.
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
pub mod room_history;
|
||||
pub mod room_info;
|
||||
pub mod script;
|
||||
pub mod search;
|
||||
|
||||
use std::sync::Arc;
|
||||
@@ -10,6 +11,7 @@ use opensearch::OpenSearch;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::context::ResponseContext;
|
||||
|
||||
pub struct ToolRegistry {
|
||||
opensearch: OpenSearch,
|
||||
@@ -122,10 +124,44 @@ impl ToolRegistry {
|
||||
"required": ["room_id"]
|
||||
}),
|
||||
),
|
||||
Tool::new(
|
||||
"run_script".into(),
|
||||
"Execute a TypeScript/JavaScript snippet in a sandboxed runtime. \
|
||||
Use this for math, date calculations, data transformations, or any \
|
||||
computation that needs precision. The script has access to:\n\
|
||||
- sol.search(query, opts?) — search the message archive. opts: \
|
||||
{ room?, sender?, after?, before?, limit?, semantic? }\n\
|
||||
- sol.rooms() — list joined rooms (returns array of {name, id, members})\n\
|
||||
- sol.members(roomName) — get room members (returns array of {name, id})\n\
|
||||
- sol.fetch(url) — HTTP GET (allowlisted domains only)\n\
|
||||
- sol.memory.get(query?) — retrieve internal notes relevant to the query\n\
|
||||
- sol.memory.set(content, category?) — save an internal note for later reference\n\
|
||||
- sol.fs.read(path), sol.fs.write(path, content), sol.fs.list(path?) — \
|
||||
sandboxed temp filesystem for intermediate files\n\
|
||||
- console.log() to produce output\n\
|
||||
All sol.* methods are async — use await. The last expression value is \
|
||||
also captured. Output is truncated to 4096 chars."
|
||||
.into(),
|
||||
json!({
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"code": {
|
||||
"type": "string",
|
||||
"description": "TypeScript or JavaScript code to execute"
|
||||
}
|
||||
},
|
||||
"required": ["code"]
|
||||
}),
|
||||
),
|
||||
]
|
||||
}
|
||||
|
||||
pub async fn execute(&self, name: &str, arguments: &str) -> anyhow::Result<String> {
|
||||
pub async fn execute(
|
||||
&self,
|
||||
name: &str,
|
||||
arguments: &str,
|
||||
response_ctx: &ResponseContext,
|
||||
) -> anyhow::Result<String> {
|
||||
match name {
|
||||
"search_archive" => {
|
||||
search::search_archive(
|
||||
@@ -145,6 +181,16 @@ impl ToolRegistry {
|
||||
}
|
||||
"list_rooms" => room_info::list_rooms(&self.matrix).await,
|
||||
"get_room_members" => room_info::get_room_members(&self.matrix, arguments).await,
|
||||
"run_script" => {
|
||||
script::run_script(
|
||||
&self.opensearch,
|
||||
&self.matrix,
|
||||
&self.config,
|
||||
arguments,
|
||||
response_ctx,
|
||||
)
|
||||
.await
|
||||
}
|
||||
_ => anyhow::bail!("Unknown tool: {name}"),
|
||||
}
|
||||
}
|
||||
|
||||
706
src/tools/script.rs
Normal file
706
src/tools/script.rs
Normal file
@@ -0,0 +1,706 @@
|
||||
use std::cell::RefCell;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::rc::Rc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use deno_core::{extension, op2, JsRuntime, OpState, RuntimeOptions};
|
||||
use deno_error::JsErrorBox;
|
||||
use matrix_sdk::Client as MatrixClient;
|
||||
use opensearch::OpenSearch;
|
||||
use serde::Deserialize;
|
||||
use tempfile::TempDir;
|
||||
use tracing::info;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::context::ResponseContext;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// State types stored in OpState
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
struct ScriptState {
|
||||
opensearch: OpenSearch,
|
||||
matrix: MatrixClient,
|
||||
config: Arc<Config>,
|
||||
tmpdir: PathBuf,
|
||||
user_id: String,
|
||||
}
|
||||
|
||||
struct ScriptOutput(String);
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct RunScriptArgs {
|
||||
code: String,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Sandbox path resolution
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn resolve_sandbox_path(
|
||||
tmpdir: &Path,
|
||||
requested: &str,
|
||||
) -> Result<PathBuf, JsErrorBox> {
|
||||
let base = tmpdir
|
||||
.canonicalize()
|
||||
.map_err(|e| JsErrorBox::generic(format!("sandbox error: {e}")))?;
|
||||
let joined = base.join(requested);
|
||||
|
||||
if let Some(parent) = joined.parent() {
|
||||
std::fs::create_dir_all(parent).ok();
|
||||
}
|
||||
|
||||
let resolved = if joined.exists() {
|
||||
joined
|
||||
.canonicalize()
|
||||
.map_err(|e| JsErrorBox::generic(format!("path error: {e}")))?
|
||||
} else {
|
||||
let parent = joined
|
||||
.parent()
|
||||
.ok_or_else(|| JsErrorBox::generic("invalid path"))?
|
||||
.canonicalize()
|
||||
.map_err(|e| JsErrorBox::generic(format!("path error: {e}")))?;
|
||||
parent.join(joined.file_name().unwrap_or_default())
|
||||
};
|
||||
|
||||
if !resolved.starts_with(&base) {
|
||||
return Err(JsErrorBox::generic("path escapes sandbox"));
|
||||
}
|
||||
|
||||
Ok(resolved)
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Ops — async (search, rooms, members, fetch)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[op2]
|
||||
#[string]
|
||||
async fn op_sol_search(
|
||||
state: Rc<RefCell<OpState>>,
|
||||
#[string] query: String,
|
||||
#[string] opts_json: String,
|
||||
) -> Result<String, JsErrorBox> {
|
||||
let (os, index) = {
|
||||
let st = state.borrow();
|
||||
let ss = st.borrow::<ScriptState>();
|
||||
(ss.opensearch.clone(), ss.config.opensearch.index.clone())
|
||||
};
|
||||
|
||||
let mut args: serde_json::Value =
|
||||
serde_json::from_str(&opts_json).unwrap_or(serde_json::json!({}));
|
||||
args["query"] = serde_json::Value::String(query);
|
||||
|
||||
super::search::search_archive(&os, &index, &args.to_string())
|
||||
.await
|
||||
.map_err(|e| JsErrorBox::generic(e.to_string()))
|
||||
}
|
||||
|
||||
#[op2]
|
||||
#[string]
|
||||
async fn op_sol_rooms(
|
||||
state: Rc<RefCell<OpState>>,
|
||||
) -> Result<String, JsErrorBox> {
|
||||
let matrix = {
|
||||
let st = state.borrow();
|
||||
st.borrow::<ScriptState>().matrix.clone()
|
||||
};
|
||||
|
||||
let rooms = matrix.joined_rooms();
|
||||
let names: Vec<serde_json::Value> = rooms
|
||||
.iter()
|
||||
.map(|r| {
|
||||
let name = match r.cached_display_name() {
|
||||
Some(n) => n.to_string(),
|
||||
None => r.room_id().to_string(),
|
||||
};
|
||||
serde_json::json!({
|
||||
"name": name,
|
||||
"id": r.room_id().to_string(),
|
||||
"members": r.joined_members_count(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
serde_json::to_string(&names).map_err(|e| JsErrorBox::generic(e.to_string()))
|
||||
}
|
||||
|
||||
#[op2]
|
||||
#[string]
|
||||
async fn op_sol_members(
|
||||
state: Rc<RefCell<OpState>>,
|
||||
#[string] room_name: String,
|
||||
) -> Result<String, JsErrorBox> {
|
||||
let matrix = {
|
||||
let st = state.borrow();
|
||||
st.borrow::<ScriptState>().matrix.clone()
|
||||
};
|
||||
|
||||
let rooms = matrix.joined_rooms();
|
||||
let room = rooms.iter().find(|r| {
|
||||
r.cached_display_name()
|
||||
.map(|n| n.to_string() == room_name)
|
||||
.unwrap_or(false)
|
||||
});
|
||||
|
||||
let Some(room) = room else {
|
||||
return Err(JsErrorBox::generic(format!(
|
||||
"room not found: {room_name}"
|
||||
)));
|
||||
};
|
||||
|
||||
let members = room
|
||||
.members(matrix_sdk::RoomMemberships::JOIN)
|
||||
.await
|
||||
.map_err(|e| JsErrorBox::generic(e.to_string()))?;
|
||||
|
||||
let names: Vec<serde_json::Value> = members
|
||||
.iter()
|
||||
.map(|m| {
|
||||
serde_json::json!({
|
||||
"name": m.display_name().unwrap_or_else(|| m.user_id().as_str()),
|
||||
"id": m.user_id().to_string(),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
serde_json::to_string(&names).map_err(|e| JsErrorBox::generic(e.to_string()))
|
||||
}
|
||||
|
||||
#[op2]
|
||||
#[string]
|
||||
async fn op_sol_fetch(
|
||||
state: Rc<RefCell<OpState>>,
|
||||
#[string] url_str: String,
|
||||
) -> Result<String, JsErrorBox> {
|
||||
let allowlist = {
|
||||
let st = state.borrow();
|
||||
st.borrow::<ScriptState>()
|
||||
.config
|
||||
.behavior
|
||||
.script_fetch_allowlist
|
||||
.clone()
|
||||
};
|
||||
|
||||
let parsed = url::Url::parse(&url_str)
|
||||
.map_err(|e| JsErrorBox::generic(format!("invalid URL: {e}")))?;
|
||||
let domain = parsed
|
||||
.host_str()
|
||||
.ok_or_else(|| JsErrorBox::generic("URL has no host"))?;
|
||||
|
||||
if !allowlist
|
||||
.iter()
|
||||
.any(|d| domain == d || domain.ends_with(&format!(".{d}")))
|
||||
{
|
||||
return Err(JsErrorBox::generic(format!(
|
||||
"domain not in allowlist: {domain}"
|
||||
)));
|
||||
}
|
||||
|
||||
let resp = reqwest::get(&url_str)
|
||||
.await
|
||||
.map_err(|e| JsErrorBox::generic(e.to_string()))?;
|
||||
|
||||
let text = resp
|
||||
.text()
|
||||
.await
|
||||
.map_err(|e| JsErrorBox::generic(e.to_string()))?;
|
||||
|
||||
let max_len = 32768;
|
||||
if text.len() > max_len {
|
||||
Ok(format!("{}...(truncated)", &text[..max_len]))
|
||||
} else {
|
||||
Ok(text)
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Ops — sync (filesystem sandbox + output collection)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[op2]
|
||||
#[string]
|
||||
fn op_sol_read_file(
|
||||
state: &mut OpState,
|
||||
#[string] path: String,
|
||||
) -> Result<String, JsErrorBox> {
|
||||
let tmpdir = state.borrow::<ScriptState>().tmpdir.clone();
|
||||
let resolved = resolve_sandbox_path(&tmpdir, &path)?;
|
||||
std::fs::read_to_string(&resolved)
|
||||
.map_err(|e| JsErrorBox::generic(format!("read error: {e}")))
|
||||
}
|
||||
|
||||
#[op2(fast)]
|
||||
fn op_sol_write_file(
|
||||
state: &mut OpState,
|
||||
#[string] path: String,
|
||||
#[string] content: String,
|
||||
) -> Result<(), JsErrorBox> {
|
||||
let tmpdir = state.borrow::<ScriptState>().tmpdir.clone();
|
||||
let resolved = resolve_sandbox_path(&tmpdir, &path)?;
|
||||
std::fs::write(&resolved, content)
|
||||
.map_err(|e| JsErrorBox::generic(format!("write error: {e}")))
|
||||
}
|
||||
|
||||
#[op2]
|
||||
#[string]
|
||||
fn op_sol_list_dir(
|
||||
state: &mut OpState,
|
||||
#[string] path: String,
|
||||
) -> Result<String, JsErrorBox> {
|
||||
let tmpdir = state.borrow::<ScriptState>().tmpdir.clone();
|
||||
let resolved = resolve_sandbox_path(&tmpdir, &path)?;
|
||||
|
||||
let entries: Vec<String> = std::fs::read_dir(&resolved)
|
||||
.map_err(|e| JsErrorBox::generic(format!("list error: {e}")))?
|
||||
.filter_map(|e| e.ok())
|
||||
.map(|e| e.file_name().to_string_lossy().into_owned())
|
||||
.collect();
|
||||
|
||||
serde_json::to_string(&entries).map_err(|e| JsErrorBox::generic(e.to_string()))
|
||||
}
|
||||
|
||||
#[op2(fast)]
|
||||
fn op_sol_set_output(
|
||||
state: &mut OpState,
|
||||
#[string] output: String,
|
||||
) -> Result<(), JsErrorBox> {
|
||||
*state.borrow_mut::<ScriptOutput>() = ScriptOutput(output);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Ops — async (memory)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[op2]
|
||||
#[string]
|
||||
async fn op_sol_memory_get(
|
||||
state: Rc<RefCell<OpState>>,
|
||||
#[string] query: String,
|
||||
) -> Result<String, JsErrorBox> {
|
||||
let (os, index, user_id) = {
|
||||
let st = state.borrow();
|
||||
let ss = st.borrow::<ScriptState>();
|
||||
(
|
||||
ss.opensearch.clone(),
|
||||
ss.config.opensearch.memory_index.clone(),
|
||||
ss.user_id.clone(),
|
||||
)
|
||||
};
|
||||
|
||||
let results = if query.is_empty() {
|
||||
crate::memory::store::get_recent(&os, &index, &user_id, 10)
|
||||
.await
|
||||
.map_err(|e| JsErrorBox::generic(e.to_string()))?
|
||||
} else {
|
||||
crate::memory::store::query(&os, &index, &user_id, &query, 10)
|
||||
.await
|
||||
.map_err(|e| JsErrorBox::generic(e.to_string()))?
|
||||
};
|
||||
|
||||
let items: Vec<serde_json::Value> = results
|
||||
.iter()
|
||||
.map(|m| {
|
||||
serde_json::json!({
|
||||
"content": m.content,
|
||||
"category": m.category,
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
|
||||
serde_json::to_string(&items).map_err(|e| JsErrorBox::generic(e.to_string()))
|
||||
}
|
||||
|
||||
#[op2]
|
||||
#[string]
|
||||
async fn op_sol_memory_set(
|
||||
state: Rc<RefCell<OpState>>,
|
||||
#[string] content: String,
|
||||
#[string] category: String,
|
||||
) -> Result<String, JsErrorBox> {
|
||||
let (os, index, user_id) = {
|
||||
let st = state.borrow();
|
||||
let ss = st.borrow::<ScriptState>();
|
||||
(
|
||||
ss.opensearch.clone(),
|
||||
ss.config.opensearch.memory_index.clone(),
|
||||
ss.user_id.clone(),
|
||||
)
|
||||
};
|
||||
|
||||
crate::memory::store::set(&os, &index, &user_id, &content, &category, "script")
|
||||
.await
|
||||
.map_err(|e| JsErrorBox::generic(e.to_string()))?;
|
||||
|
||||
Ok("ok".into())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
extension!(
|
||||
sol_script,
|
||||
ops = [
|
||||
op_sol_search,
|
||||
op_sol_rooms,
|
||||
op_sol_members,
|
||||
op_sol_fetch,
|
||||
op_sol_read_file,
|
||||
op_sol_write_file,
|
||||
op_sol_list_dir,
|
||||
op_sol_set_output,
|
||||
op_sol_memory_get,
|
||||
op_sol_memory_set,
|
||||
],
|
||||
state = |state| {
|
||||
state.put(ScriptOutput(String::new()));
|
||||
},
|
||||
);
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Bootstrap JS — injected before user code
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
const BOOTSTRAP_JS: &str = r#"
|
||||
const __output = [];
|
||||
globalThis.console = {
|
||||
log: (...args) => __output.push(args.map(a => typeof a === 'object' ? JSON.stringify(a) : String(a)).join(' ')),
|
||||
error: (...args) => __output.push('ERROR: ' + args.map(a => typeof a === 'object' ? JSON.stringify(a) : String(a)).join(' ')),
|
||||
warn: (...args) => __output.push('WARN: ' + args.map(a => typeof a === 'object' ? JSON.stringify(a) : String(a)).join(' ')),
|
||||
info: (...args) => __output.push(args.map(a => typeof a === 'object' ? JSON.stringify(a) : String(a)).join(' ')),
|
||||
};
|
||||
globalThis.sol = {
|
||||
search: (query, opts) => Deno.core.ops.op_sol_search(query, opts ? JSON.stringify(opts) : "{}"),
|
||||
rooms: async () => JSON.parse(await Deno.core.ops.op_sol_rooms()),
|
||||
members: async (room) => JSON.parse(await Deno.core.ops.op_sol_members(room)),
|
||||
fetch: (url) => Deno.core.ops.op_sol_fetch(url),
|
||||
memory: {
|
||||
get: async (query) => JSON.parse(await Deno.core.ops.op_sol_memory_get(query || "")),
|
||||
set: async (content, category) => {
|
||||
await Deno.core.ops.op_sol_memory_set(content, category || "general");
|
||||
},
|
||||
},
|
||||
fs: {
|
||||
read: (path) => Deno.core.ops.op_sol_read_file(path),
|
||||
write: (path, content) => Deno.core.ops.op_sol_write_file(path, content),
|
||||
list: (path) => JSON.parse(Deno.core.ops.op_sol_list_dir(path || ".")),
|
||||
},
|
||||
};
|
||||
"#;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// TypeScript transpilation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn transpile_ts(code: &str) -> anyhow::Result<String> {
|
||||
use deno_ast::{MediaType, ParseParams, TranspileModuleOptions};
|
||||
|
||||
let specifier = deno_ast::ModuleSpecifier::parse("file:///script.ts")?;
|
||||
|
||||
let parsed = deno_ast::parse_module(ParseParams {
|
||||
specifier,
|
||||
text: code.into(),
|
||||
media_type: MediaType::TypeScript,
|
||||
capture_tokens: false,
|
||||
maybe_syntax: None,
|
||||
scope_analysis: false,
|
||||
})?;
|
||||
|
||||
let transpiled = parsed.transpile(
|
||||
&deno_ast::TranspileOptions::default(),
|
||||
&TranspileModuleOptions::default(),
|
||||
&deno_ast::EmitOptions::default(),
|
||||
)?;
|
||||
|
||||
Ok(transpiled.into_source().text.to_string())
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Public entry point
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
pub async fn run_script(
|
||||
opensearch: &OpenSearch,
|
||||
matrix: &MatrixClient,
|
||||
config: &Config,
|
||||
args_json: &str,
|
||||
response_ctx: &ResponseContext,
|
||||
) -> anyhow::Result<String> {
|
||||
let args: RunScriptArgs = serde_json::from_str(args_json)?;
|
||||
let code = args.code.clone();
|
||||
|
||||
info!(code_len = code.len(), "Executing script");
|
||||
|
||||
// Transpile TS to JS
|
||||
let js_code = transpile_ts(&code)?;
|
||||
|
||||
// Clone state for move into spawn_blocking
|
||||
let os = opensearch.clone();
|
||||
let mx = matrix.clone();
|
||||
let cfg = Arc::new(config.clone());
|
||||
let timeout_secs = cfg.behavior.script_timeout_secs;
|
||||
let max_heap_mb = cfg.behavior.script_max_heap_mb;
|
||||
let user_id = response_ctx.user_id.clone();
|
||||
|
||||
// Wrap user code: async IIFE that captures output, then stores via op
|
||||
let wrapped = format!(
|
||||
r#"
|
||||
(async () => {{
|
||||
try {{
|
||||
const __result = await (async () => {{
|
||||
{js_code}
|
||||
}})();
|
||||
if (__result !== undefined) {{
|
||||
__output.push(typeof __result === 'object' ? JSON.stringify(__result, null, 2) : String(__result));
|
||||
}}
|
||||
}} catch(e) {{
|
||||
__output.push('Error: ' + (e.stack || e.message || String(e)));
|
||||
}}
|
||||
Deno.core.ops.op_sol_set_output(JSON.stringify(__output));
|
||||
}})();"#
|
||||
);
|
||||
|
||||
let timeout = Duration::from_secs(timeout_secs);
|
||||
|
||||
let result = tokio::task::spawn_blocking(move || {
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()?;
|
||||
|
||||
rt.block_on(async move {
|
||||
let tmpdir = TempDir::new()?;
|
||||
let tmpdir_path = tmpdir.path().to_path_buf();
|
||||
|
||||
let mut runtime = JsRuntime::new(RuntimeOptions {
|
||||
extensions: vec![sol_script::init()],
|
||||
create_params: Some(
|
||||
deno_core::v8::CreateParams::default()
|
||||
.heap_limits(0, max_heap_mb * 1024 * 1024),
|
||||
),
|
||||
..Default::default()
|
||||
});
|
||||
|
||||
// Inject shared state
|
||||
{
|
||||
let op_state = runtime.op_state();
|
||||
let mut state = op_state.borrow_mut();
|
||||
state.put(ScriptState {
|
||||
opensearch: os,
|
||||
matrix: mx,
|
||||
config: cfg,
|
||||
tmpdir: tmpdir_path,
|
||||
user_id,
|
||||
});
|
||||
}
|
||||
|
||||
// V8 isolate termination for timeout
|
||||
let done = Arc::new(AtomicBool::new(false));
|
||||
let done_clone = done.clone();
|
||||
let isolate_handle = runtime.v8_isolate().thread_safe_handle();
|
||||
std::thread::spawn(move || {
|
||||
let deadline = std::time::Instant::now() + timeout;
|
||||
while std::time::Instant::now() < deadline {
|
||||
if done_clone.load(Ordering::Relaxed) {
|
||||
return;
|
||||
}
|
||||
std::thread::sleep(Duration::from_millis(50));
|
||||
}
|
||||
isolate_handle.terminate_execution();
|
||||
});
|
||||
|
||||
// Execute bootstrap
|
||||
runtime.execute_script("<bootstrap>", BOOTSTRAP_JS)?;
|
||||
|
||||
// Execute user code
|
||||
let exec_result = runtime.execute_script("<script>", wrapped);
|
||||
|
||||
if let Err(e) = exec_result {
|
||||
done.store(true, Ordering::Relaxed);
|
||||
let msg = e.to_string();
|
||||
if msg.contains("terminated") {
|
||||
return Ok("Error: script execution timed out".into());
|
||||
}
|
||||
return Ok(format!("Error: {msg}"));
|
||||
}
|
||||
|
||||
// Drive async ops to completion
|
||||
if let Err(e) = runtime.run_event_loop(Default::default()).await {
|
||||
done.store(true, Ordering::Relaxed);
|
||||
let msg = e.to_string();
|
||||
if msg.contains("terminated") {
|
||||
return Ok("Error: script execution timed out".into());
|
||||
}
|
||||
return Ok(format!("Error: {msg}"));
|
||||
}
|
||||
|
||||
done.store(true, Ordering::Relaxed);
|
||||
|
||||
// Read captured output from OpState
|
||||
let output_json = {
|
||||
let op_state = runtime.op_state();
|
||||
let state = op_state.borrow();
|
||||
state.borrow::<ScriptOutput>().0.clone()
|
||||
};
|
||||
|
||||
let output_arr: Vec<String> =
|
||||
serde_json::from_str(&output_json).unwrap_or_else(|_| {
|
||||
if output_json.is_empty() {
|
||||
vec![]
|
||||
} else {
|
||||
vec![output_json]
|
||||
}
|
||||
});
|
||||
let result = output_arr.join("\n");
|
||||
|
||||
let max_output = 4096;
|
||||
let result = if result.len() > max_output {
|
||||
format!("{}...(truncated)", &result[..max_output])
|
||||
} else {
|
||||
result
|
||||
};
|
||||
|
||||
drop(tmpdir);
|
||||
|
||||
Ok::<String, anyhow::Error>(result)
|
||||
})
|
||||
})
|
||||
.await?;
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_run_script_args() {
|
||||
let args: RunScriptArgs =
|
||||
serde_json::from_str(r#"{"code": "console.log(42)"}"#).unwrap();
|
||||
assert_eq!(args.code, "console.log(42)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transpile_ts_basic() {
|
||||
let js = transpile_ts("const x: number = 42; x;").unwrap();
|
||||
assert!(js.contains("const x = 42"));
|
||||
assert!(!js.contains(": number"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transpile_ts_arrow() {
|
||||
let js =
|
||||
transpile_ts("const add = (a: number, b: number): number => a + b;").unwrap();
|
||||
assert!(js.contains("const add"));
|
||||
assert!(!js.contains(": number"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transpile_ts_interface() {
|
||||
let js = transpile_ts(
|
||||
"interface Foo { bar: string; } const x: Foo = { bar: 'hello' };",
|
||||
)
|
||||
.unwrap();
|
||||
assert!(!js.contains("interface"));
|
||||
assert!(js.contains("const x"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_transpile_ts_invalid() {
|
||||
let result = transpile_ts("const x: number = {");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sandbox_path_within() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let base = dir.path();
|
||||
std::fs::write(base.join("test.txt"), "hello").unwrap();
|
||||
|
||||
let resolved = resolve_sandbox_path(base, "test.txt").unwrap();
|
||||
assert!(resolved.starts_with(base.canonicalize().unwrap()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sandbox_path_escape_rejected() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let base = dir.path();
|
||||
|
||||
// Path traversal is blocked — either by canonicalization failure
|
||||
// (parent doesn't exist) or by the starts_with sandbox check
|
||||
let result = resolve_sandbox_path(base, "../../../etc/passwd");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sandbox_path_symlink_escape_rejected() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let base = dir.path();
|
||||
|
||||
// Create a symlink pointing outside the sandbox
|
||||
let link_path = base.join("escape");
|
||||
std::os::unix::fs::symlink("/tmp", &link_path).unwrap();
|
||||
|
||||
// Following the symlink resolves outside sandbox
|
||||
let result = resolve_sandbox_path(base, "escape/somefile");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sandbox_path_new_file() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let base = dir.path();
|
||||
|
||||
let resolved = resolve_sandbox_path(base, "newfile.txt").unwrap();
|
||||
assert!(resolved.starts_with(base.canonicalize().unwrap()));
|
||||
assert!(resolved.ends_with("newfile.txt"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sandbox_path_nested_dir() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let base = dir.path();
|
||||
|
||||
let resolved = resolve_sandbox_path(base, "subdir/file.txt").unwrap();
|
||||
assert!(resolved.starts_with(base.canonicalize().unwrap()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_basic_script_execution() {
|
||||
let mut runtime = JsRuntime::new(RuntimeOptions {
|
||||
extensions: vec![sol_script::init()],
|
||||
..Default::default()
|
||||
});
|
||||
|
||||
runtime
|
||||
.execute_script("<bootstrap>", BOOTSTRAP_JS)
|
||||
.unwrap();
|
||||
|
||||
runtime
|
||||
.execute_script(
|
||||
"<test>",
|
||||
r#"
|
||||
console.log(2 ** 64);
|
||||
console.log(Math.PI);
|
||||
Deno.core.ops.op_sol_set_output(JSON.stringify(__output));
|
||||
"#,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let op_state = runtime.op_state();
|
||||
let state = op_state.borrow();
|
||||
let output = &state.borrow::<ScriptOutput>().0;
|
||||
let arr: Vec<String> = serde_json::from_str(output).unwrap();
|
||||
assert_eq!(arr.len(), 2);
|
||||
assert!(arr[0].contains("18446744073709552000"));
|
||||
assert!(arr[1].contains("3.14159"));
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,7 @@
|
||||
use opensearch::OpenSearch;
|
||||
use serde::Deserialize;
|
||||
use serde_json::json;
|
||||
use tracing::debug;
|
||||
use tracing::{debug, info};
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct SearchArgs {
|
||||
@@ -24,19 +24,24 @@ fn default_limit() -> usize { 10 }
|
||||
|
||||
/// Build the OpenSearch query body from parsed SearchArgs. Extracted for testability.
|
||||
pub fn build_search_query(args: &SearchArgs) -> serde_json::Value {
|
||||
let must = vec![json!({
|
||||
"match": { "content": args.query }
|
||||
})];
|
||||
// Handle empty/wildcard queries as match_all
|
||||
let must = if args.query.is_empty() || args.query == "*" {
|
||||
vec![json!({ "match_all": {} })]
|
||||
} else {
|
||||
vec![json!({
|
||||
"match": { "content": args.query }
|
||||
})]
|
||||
};
|
||||
|
||||
let mut filter = vec![json!({
|
||||
"term": { "redacted": false }
|
||||
})];
|
||||
|
||||
if let Some(ref room) = args.room {
|
||||
filter.push(json!({ "term": { "room_name": room } }));
|
||||
filter.push(json!({ "term": { "room_name.keyword": room } }));
|
||||
}
|
||||
if let Some(ref sender) = args.sender {
|
||||
filter.push(json!({ "term": { "sender_name": sender } }));
|
||||
filter.push(json!({ "term": { "sender_name.keyword": sender } }));
|
||||
}
|
||||
|
||||
let mut range = serde_json::Map::new();
|
||||
@@ -73,10 +78,19 @@ pub async fn search_archive(
|
||||
args_json: &str,
|
||||
) -> anyhow::Result<String> {
|
||||
let args: SearchArgs = serde_json::from_str(args_json)?;
|
||||
debug!(query = args.query.as_str(), "Searching archive");
|
||||
|
||||
let query_body = build_search_query(&args);
|
||||
|
||||
info!(
|
||||
query = args.query.as_str(),
|
||||
room = args.room.as_deref().unwrap_or("*"),
|
||||
sender = args.sender.as_deref().unwrap_or("*"),
|
||||
after = args.after.as_deref().unwrap_or("*"),
|
||||
before = args.before.as_deref().unwrap_or("*"),
|
||||
limit = args.limit,
|
||||
query_json = %query_body,
|
||||
"Executing search"
|
||||
);
|
||||
|
||||
let response = client
|
||||
.search(opensearch::SearchParts::Index(&[index]))
|
||||
.body(query_body)
|
||||
@@ -84,6 +98,8 @@ pub async fn search_archive(
|
||||
.await?;
|
||||
|
||||
let body: serde_json::Value = response.json().await?;
|
||||
let hit_count = body["hits"]["total"]["value"].as_i64().unwrap_or(0);
|
||||
info!(hit_count, "Search results");
|
||||
let hits = &body["hits"]["hits"];
|
||||
|
||||
let Some(hits_arr) = hits.as_array() else {
|
||||
@@ -173,7 +189,7 @@ mod tests {
|
||||
|
||||
let filters = q["query"]["bool"]["filter"].as_array().unwrap();
|
||||
assert_eq!(filters.len(), 2);
|
||||
assert_eq!(filters[1]["term"]["room_name"], "design");
|
||||
assert_eq!(filters[1]["term"]["room_name.keyword"], "design");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -183,7 +199,7 @@ mod tests {
|
||||
|
||||
let filters = q["query"]["bool"]["filter"].as_array().unwrap();
|
||||
assert_eq!(filters.len(), 2);
|
||||
assert_eq!(filters[1]["term"]["sender_name"], "Bob");
|
||||
assert_eq!(filters[1]["term"]["sender_name.keyword"], "Bob");
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -193,8 +209,8 @@ mod tests {
|
||||
|
||||
let filters = q["query"]["bool"]["filter"].as_array().unwrap();
|
||||
assert_eq!(filters.len(), 3);
|
||||
assert_eq!(filters[1]["term"]["room_name"], "dev");
|
||||
assert_eq!(filters[2]["term"]["sender_name"], "Carol");
|
||||
assert_eq!(filters[1]["term"]["room_name.keyword"], "dev");
|
||||
assert_eq!(filters[2]["term"]["sender_name.keyword"], "Carol");
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user