From b1f21b8f54846d29fde12edb612cd3618c707142 Mon Sep 17 00:00:00 2001 From: Edgars Date: Sun, 27 Sep 2026 17:57:22 +0100 Subject: [PATCH 1/2] =?UTF-8?q?chore(modules):=20expose=20wrapped=20budget?= =?UTF-8?q?=20timeout=20=E2=9C=85?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- implementation/src/llm/handler.rs | 4 +++ implementation/src/llm/handler_test.rs | 42 ++++++++++++++++++++++++++ 2 files changed, 46 insertions(+) create mode 100644 implementation/src/llm/handler_test.rs diff --git a/implementation/src/llm/handler.rs b/implementation/src/llm/handler.rs index 76b843a..ae80a79 100644 --- a/implementation/src/llm/handler.rs +++ b/implementation/src/llm/handler.rs @@ -20,6 +20,10 @@ pub struct Inner { genvm_id: genvm_modules_interfaces::GenVMId, } +#[cfg(test)] +#[path = "handler_test.rs"] +mod tests; + type LlmSubContext = crate::manager::execution_context::LlmSubContext; pub struct Provider { diff --git a/implementation/src/llm/handler_test.rs b/implementation/src/llm/handler_test.rs new file mode 100644 index 0000000..8138cb2 --- /dev/null +++ b/implementation/src/llm/handler_test.rs @@ -0,0 +1,42 @@ +use super::Inner; + +#[tokio::test] +async fn budget_exhaustion_from_lua_callback_becomes_timeout_sentinel() -> anyhow::Result<()> { + let direct = Inner::try_catch_budget_exhausted(anyhow::Error::from(mlua::Error::external( + crate::common::BudgetExhausted, + )))?; + assert_eq!(direct.consumed_gen, primitive_types::U256::MAX); + + let lua = mlua::Lua::new(); + lua.globals().set( + "exhaust", + lua.create_async_function(|_, ()| async { + Err::<(), mlua::Error>(mlua::Error::external(crate::common::BudgetExhausted)) + })?, + )?; + + let call = lua + .load("return function() exhaust() end") + .eval::()?; + let err = call + .call_async::<()>(()) + .await + .expect_err("budget must exhaust"); + assert!( + matches!( + err, + mlua::Error::CallbackError { ref cause, .. } + if matches!(cause.as_ref(), mlua::Error::ExternalError(ext) + if ext.downcast_ref::().is_some()) + ), + "expected Lua callback to wrap BudgetExhausted, got {err:?}" + ); + + let answer = Inner::try_catch_budget_exhausted(anyhow::Error::from(err))?; + assert_eq!(answer.consumed_gen, primitive_types::U256::MAX); + assert!(matches!( + answer.data, + genvm_modules_interfaces::llm::PromptAnswerData::Text(ref text) if text.is_empty() + )); + Ok(()) +} From cc35474cbc8055554bd1a4b96f9d2ed0c2c02099 Mon Sep 17 00:00:00 2001 From: kp2pml30 Date: Tue, 29 Sep 2026 13:47:44 +0900 Subject: [PATCH 2/2] =?UTF-8?q?fix(modules):=20catch=20budget=20exhaustion?= =?UTF-8?q?=20under=20any=20Lua=20error=20wrapper=20=F0=9F=90=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Walk the whole error chain for `BudgetExhausted`, including `CallbackError`, `BadArgument`, `WithContext` and the `Arc`s `scripting::call_fn` unwraps to - Cover a `pcall` rethrow, both `call_fn` shapes and an unrelated callback error --- implementation/src/llm/handler.rs | 43 +++++++++-------- implementation/src/llm/handler_test.rs | 65 ++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 20 deletions(-) diff --git a/implementation/src/llm/handler.rs b/implementation/src/llm/handler.rs index ae80a79..eb599f1 100644 --- a/implementation/src/llm/handler.rs +++ b/implementation/src/llm/handler.rs @@ -184,28 +184,31 @@ impl Inner { Ok(llm_iface::PromptAnswer { data, consumed_gen }) } - fn try_catch_budget_exhausted(err: anyhow::Error) -> ModuleResult { - if err - .downcast_ref::() - .is_some() - { - return Ok(llm_iface::PromptAnswer { - data: llm_iface::PromptAnswerData::Text(String::new()), - consumed_gen: primitive_types::U256::MAX, - }); + fn is_budget_exhausted(err: &(dyn std::error::Error + 'static)) -> bool { + // `scripting::call_fn` unwraps top-level `ExternalError`/`WithContext` into an `Arc` + if let Some(ext) = err.downcast_ref::>() { + return Self::is_budget_exhausted(&**ext); } - if let Some(mlua::Error::ExternalError(ext)) = err.downcast_ref::() { - if ext - .downcast_ref::() - .is_some() - { - return Ok(llm_iface::PromptAnswer { - data: llm_iface::PromptAnswerData::Text(String::new()), - consumed_gen: primitive_types::U256::MAX, - }); - } + if let Some(cause) = err.downcast_ref::>() { + return Self::is_budget_exhausted(&**cause); + } + match err.downcast_ref::() { + // walks `CallbackError`, `BadArgument`, `WithContext` and `ExternalError` wrappers + Some(lua_err) => lua_err + .chain() + .any(|e| e.is::()), + None => err.is::(), + } + } + + fn try_catch_budget_exhausted(err: anyhow::Error) -> ModuleResult { + if !err.chain().any(Self::is_budget_exhausted) { + return Err(err); } - Err(err) + Ok(llm_iface::PromptAnswer { + data: llm_iface::PromptAnswerData::Text(String::new()), + consumed_gen: primitive_types::U256::MAX, + }) } async fn exec_prompt( diff --git a/implementation/src/llm/handler_test.rs b/implementation/src/llm/handler_test.rs index 8138cb2..b3928d8 100644 --- a/implementation/src/llm/handler_test.rs +++ b/implementation/src/llm/handler_test.rs @@ -40,3 +40,68 @@ async fn budget_exhaustion_from_lua_callback_becomes_timeout_sentinel() -> anyho )); Ok(()) } + +#[tokio::test] +async fn budget_exhaustion_survives_lua_rethrow() -> anyhow::Result<()> { + let lua = mlua::Lua::new(); + lua.globals().set( + "exhaust", + lua.create_function(|_, ()| -> mlua::Result<()> { + Err(mlua::Error::external(crate::common::BudgetExhausted)) + })?, + )?; + + let call = lua + .load( + r#" + local function inner() exhaust() end + return function() + local ok, err = pcall(inner) + if not ok then error(err, 0) end + end + "#, + ) + .eval::()?; + let err = call + .call_async::<()>(()) + .await + .expect_err("budget must exhaust"); + + let answer = Inner::try_catch_budget_exhausted(anyhow::Error::from(err))?; + assert_eq!(answer.consumed_gen, primitive_types::U256::MAX); + Ok(()) +} + +#[test] +fn budget_exhaustion_unwrapped_by_call_fn_becomes_timeout_sentinel() -> anyhow::Result<()> { + let ext: std::sync::Arc = + std::sync::Arc::new(crate::common::BudgetExhausted); + let answer = Inner::try_catch_budget_exhausted(anyhow::Error::from(ext))?; + assert_eq!(answer.consumed_gen, primitive_types::U256::MAX); + + let cause = std::sync::Arc::new(mlua::Error::CallbackError { + traceback: String::new(), + cause: std::sync::Arc::new(mlua::Error::external(crate::common::BudgetExhausted)), + }); + let answer = + Inner::try_catch_budget_exhausted(anyhow::Error::from(cause).context("exec_prompt"))?; + assert_eq!(answer.consumed_gen, primitive_types::U256::MAX); + Ok(()) +} + +#[tokio::test] +async fn unrelated_lua_callback_error_is_not_budget_exhaustion() -> anyhow::Result<()> { + let lua = mlua::Lua::new(); + lua.globals().set( + "fail", + lua.create_function(|_, ()| -> mlua::Result<()> { Err(mlua::Error::runtime("boom")) })?, + )?; + + let call = lua + .load("return function() fail() end") + .eval::()?; + let err = call.call_async::<()>(()).await.expect_err("must fail"); + + assert!(Inner::try_catch_budget_exhausted(anyhow::Error::from(err)).is_err()); + Ok(()) +}