@layerzerolabs/framework-solana-anchor-trait 0.2.62 → 0.2.64

package/src/generation.rs

@@ -1,60 +1,97 @@
 //! Instruction code generation with override level resolution.
 //!
 //! - **Default**: No override → emit default impl + default entrypoint.
-//! - **Handler Override**: Override with default context type → emit default impl, entrypoint
-//!   dispatches to implementor's handler.
-//! - **Full Override**: Override with custom context type → skip default impl, entrypoint
-//!   dispatches to implementor's handler with custom type.
+//! - **Override**: Custom handler present → emit default impl (always available), entrypoint
+//!   dispatches to implementor's handler using whatever context type the handler declares.
 
 use proc_macro2::{Ident, Span, TokenStream};
 use quote::{quote, ToTokens};
+use syn::{ItemImpl, ItemStruct};
 
-use crate::overrides::{HandlerMeta, InstructionOverrides};
+use crate::overrides::InstructionOverrides;
 use crate::InstructionSet;
 
 /// Descriptor for one instruction's code generation.
 pub struct InstructionSpec<Ctx> {
-    /// Parameter type as a token stream (e.g. `Some(quote!(::oapp::params::SetPeerParams))`).
+    /// Parameter type as a token stream (e.g. `Some(quote!(::lz_oapp::types::SetPeerParams))`).
     /// Use `None` for instructions with no params.
     pub params_type: Option<TokenStream>,
     /// Return type as a token stream (e.g. `quote!(())` or `quote!([u8; 32])`).
     pub return_type: TokenStream,
-    /// When `true`, the entrypoint passes `ctx` / `&ctx` (immutable).
-    /// When `false` , it passes `mut ctx` / `&mut ctx`.
+    /// When `true`, the entrypoint passes `&ctx` (immutable).
+    /// When `false`, it passes `&mut ctx`.
     pub is_view: bool,
-    /// - `Some(...)` — provided instruction (all three override levels valid).
-    /// - `None` — **required** instruction: the implementor must provide an override handler (only
-    ///   Full Override is valid). Omitting the override produces a compile error.
+    /// - `Some(...)` — provided instruction (Default and Override levels valid).
+    /// - `None` — **required** instruction: the implementor must provide an override handler.
+    ///   Omitting the override produces a compile error.
     pub default_impl: Option<DefaultImpl<Ctx>>,
 }
 
 /// Default impl generation for an instruction (accounts struct + handler).
 ///
-/// Groups the three fields that are only meaningful when an instruction
+/// Groups the fields that are only meaningful when an instruction
 /// has a framework-provided default implementation.
 pub struct DefaultImpl<Ctx> {
-    /// Name of the default accounts struct, e.g. `format_ident!("StdSetPeer")`.
-    /// Also used to detect Handler Override vs Full Override (string comparison with handler's
-    /// accounts type).
-    pub accounts_ident: Ident,
     /// Name of the default handler method on the context struct, e.g. `format_ident!("apply")`.
     /// At the Default level, the entrypoint calls `StdAccounts::<this>(...)`.
     pub handler_ident: Ident,
-    /// Generates the default accounts struct + handler impl.
-    /// The impl must define a method matching [`handler_ident`](Self::handler_ident).
-    /// For mutable instructions: `fn(&mut Context<Self>, &Params) -> Result<R>`.
-    /// For view instructions: `fn(&Context<Self>, &Params) -> Result<R>`.
-    /// Omit the params argument when `params_type = None`.
-    pub generator: fn(&Ctx, &InstructionSpec<Ctx>) -> TokenStream,
+    /// Returns the complete accounts struct and its impl block as typed AST.
+    pub generator: fn(&Ctx, &InstructionSpec<Ctx>) -> DefaultAccounts,
 }
 
-impl<Ctx> InstructionSpec<Ctx> {
-    /// Returns `true` if the given context's accounts type matches the default accounts ident,
-    /// meaning the override handler reuses the default context struct (Handler Override).
-    fn uses_default(&self, accounts_type: &syn::Type) -> bool {
-        self.default_impl.as_ref().is_some_and(|default| {
-            matches!(accounts_type, syn::Type::Path(ty) if ty.path.is_ident(&default.accounts_ident))
-        })
+/// Complete accounts struct and its handler impl block.
+///
+/// The framework uses this template to produce two outputs:
+/// 1. The concrete struct + impl — emitted as-is via [`to_definition`](Self::to_definition).
+/// 2. A `macro_rules!` rule — built by decomposing the struct into attrs/fields and replacing the
+///    name with `$name` via [`to_extend_rule`](Self::to_extend_rule).
+pub struct DefaultAccounts {
+    /// Complete struct definition including attributes, name, generics, and fields.
+    pub struct_type: ItemStruct,
+    /// Complete impl block with handler methods. Must use `Context<Self>` so the
+    /// methods work regardless of the struct name.
+    pub impl_block: ItemImpl,
+}
+
+impl DefaultAccounts {
+    /// Returns the struct ident (e.g. `StdSetPeer`) — the single source of truth.
+    pub fn accounts_ident(&self) -> &Ident {
+        &self.struct_type.ident
+    }
+
+    /// Emits the struct and impl block.
+    pub fn to_definition(&self) -> TokenStream {
+        let s = &self.struct_type;
+        let i = &self.impl_block;
+        quote! { #s #i }
+    }
+
+    /// Assembles one `macro_rules!` rule for extending the default accounts.
+    ///
+    /// Decomposes the struct into attributes and fields, then re-emits with
+    /// `$name` replacing the struct ident and `$($extra)*` appended to the fields.
+    pub fn to_extend_rule(&self, instruction_name: &str) -> TokenStream {
+        let attrs = &self.struct_type.attrs;
+        let fields: Vec<_> = match &self.struct_type.fields {
+            syn::Fields::Named(f) => f.named.iter().collect(),
+            _ => unreachable!("accounts structs always have named fields"),
+        };
+        let impl_items = &self.impl_block.items;
+        let rule_ident = Ident::new(instruction_name, Span::call_site());
+
+        quote! {
+            (#rule_ident, $name:ident, { $($extra:tt)* }) => {
+                #(#attrs)*
+                pub struct $name<'info> {
+                    #(#fields,)*
+                    $($extra)*
+                }
+
+                impl $name<'_> {
+                    #(#impl_items)*
+                }
+            };
+        }
     }
 }
 
@@ -66,31 +103,18 @@ pub struct GeneratedInstructions {
     pub default_impls: TokenStream,
     /// All entrypoint function definitions (placed inside the program module).
     pub entrypoints: TokenStream,
+    /// Collected `macro_rules!` rules for extending default accounts structs.
+    pub extend_rules: TokenStream,
 }
 
 impl GeneratedInstructions {
     fn extend(&mut self, other: Self) {
         self.default_impls.extend(other.default_impls);
         self.entrypoints.extend(other.entrypoints);
+        self.extend_rules.extend(other.extend_rules);
     }
 }
 
-/// Resolved override: everything needed to assemble the entrypoint.
-struct ResolvedOverride {
-    /// Tokens for the default accounts struct + handler impl (emitted outside the
-    /// program module). Empty when the override fully replaces the default context.
-    default_impl: TokenStream,
-    /// Function-level generics from the override handler (e.g. `<'info>`).
-    /// Empty at the Default level since the generated entrypoint needs no extra generics.
-    fn_generics: TokenStream,
-    /// The `ctx` parameter declaration: `ctx: Context<T>` or `mut ctx: Context<T>`.
-    ctx_param: TokenStream,
-    /// Body expression for instructions that take params: `handler(ctx, params)`.
-    call_with_params: TokenStream,
-    /// Body expression for parameter-less instructions: `handler(ctx)`.
-    call_no_params: TokenStream,
-}
-
 /// Generates code for all instructions and merges into a single output.
 pub fn generate_instructions<Ix: InstructionSet>(
     ctx: &Ix::Ctx,
@@ -103,7 +127,7 @@ pub fn generate_instructions<Ix: InstructionSet>(
     Ok(out)
 }
 
-/// Generates code for a single instruction with override level resolution.
+/// Generates code for a single instruction.
 ///
 /// Returns `Err` if the instruction is required (`default_impl = None`)
 /// and the implementor did not provide an override handler.
@@ -117,110 +141,63 @@ fn generate_instruction<Ix: InstructionSet>(
 
     // Required instruction with no override → compile error.
     if handler.is_none() && spec.default_impl.is_none() {
-        let cfg = &overrides.config;
         return Err(syn::Error::new(
             Span::call_site(),
             format!(
                 "instruction `{}` has no default implementation; \
-                 provide a handler with #[{}({})]",
+                 provide a `fn {}(...)` annotated with #[{}]",
                 ix_name.canonical_name(),
-                cfg.attr_name,
                 ix_name.canonical_name(),
+                &overrides.config.attr_name,
             ),
         ));
     }
 
-    let fn_ident = syn::Ident::new(ix_name.canonical_name(), Span::call_site());
-    let return_type = &spec.return_type;
+    // Always emit the default accounts when one exists (available even under override).
+    // Also collect the extend rule for the macro_rules! arm.
+    let default_accounts = spec.default_impl.as_ref().map(|d| (d.generator)(ctx, &spec));
+    let (default_impls, extend_rules) = match &default_accounts {
+        Some(da) => (da.to_definition(), da.to_extend_rule(ix_name.canonical_name())),
+        None => (TokenStream::new(), TokenStream::new()),
+    };
 
-    let ResolvedOverride { default_impl, fn_generics, ctx_param, call_with_params, call_no_params } =
-        resolve_override(handler, &spec, ctx);
+    // Assemble entrypoint fragments from spec + dispatch target.
+    let entry_fn_ident = syn::Ident::new(ix_name.canonical_name(), Span::call_site());
 
-    // Assemble the entrypoint function, with or without a params argument.
-    // Override handlers preserve function generics (e.g. lifetime params for remaining_accounts).
-    let entrypoints = match &spec.params_type {
-        Some(params_type) => quote! {
-            pub fn #fn_ident #fn_generics (#ctx_param, params: #params_type) -> Result<#return_type> {
-                #call_with_params
-            }
+    // Resolve dispatch target: override handler vs default handler.
+    let (fn_generics, context_type, dispatch) = match handler {
+        Some(h) => {
+            let wrapper = &h.wrapper_mod;
+            let ctx_type = &h.context_type;
+            (h.fn_generics.to_token_stream(), quote!(#ctx_type), quote!(#wrapper::#entry_fn_ident))
         },
-        None => quote! {
-            pub fn #fn_ident #fn_generics (#ctx_param) -> Result<#return_type> {
-                #call_no_params
-            }
+        None => {
+            let default = spec.default_impl.as_ref().expect("validated above");
+            let accounts = default_accounts.as_ref().unwrap().accounts_ident();
+            let handler_fn = &default.handler_ident;
+            (TokenStream::new(), quote!(Context<#accounts>), quote!(#accounts::#handler_fn))
         },
     };
+    let return_type = &spec.return_type;
+    let (ref_token, ctx_decl) = if spec.is_view {
+        (quote!(&), quote!(ctx: #context_type))
+    } else {
+        (quote!(&mut), quote!(mut ctx: #context_type))
+    };
 
-    Ok(GeneratedInstructions { default_impls: default_impl, entrypoints })
-}
-
-/// Resolves which override level applies and returns the tokens for each part.
-///
-/// - **Override present** → dispatch to the implementor's wrapper handler. Emit default impl only
-///   if the handler reuses the default context (Handler Override).
-/// - **No override** → dispatch to the default handler directly.
-fn resolve_override<Ctx>(
-    override_handler: Option<&HandlerMeta>,
-    spec: &InstructionSpec<Ctx>,
-    ctx: &Ctx,
-) -> ResolvedOverride {
-    match (override_handler, &spec.default_impl) {
-        // Handler Override or Full Override — the implementor provided a custom handler.
-        //
-        // Handler Override (context type matches default): emit default impl so the
-        //   override handler can delegate back to the default accounts struct.
-        // Full Override (different context type): skip default impl entirely; the
-        //   implementor supplies their own accounts struct and handler logic.
-        //
-        // In both cases the entrypoint dispatches to the implementor's wrapper handler,
-        // passing `ctx` and `params` by value (ownership) so they can freely mutate,
-        // destructure, or forward them.
-        (Some(handler), spec_default) => {
-            let default_impl = if spec.uses_default(&handler.accounts_type) {
-                (spec_default.as_ref().unwrap().generator)(ctx, spec)
-            } else {
-                TokenStream::new()
-            };
-            let wrapper = &handler.wrapper_mod;
-            let handler_fn = &handler.fn_ident;
-            let context_type = &handler.context_type;
-            ResolvedOverride {
-                default_impl,
-                fn_generics: handler.fn_generics.to_token_stream(),
-                ctx_param: ctx_param_tokens(spec.is_view, quote!(#context_type)),
-                call_with_params: quote!(#wrapper::#handler_fn(ctx, params)),
-                call_no_params: quote!(#wrapper::#handler_fn(ctx)),
+    // Emit the entrypoint function, with or without a params argument.
+    let entrypoints = match &spec.params_type {
+        Some(params_type) => quote! {
+            pub fn #entry_fn_ident #fn_generics (#ctx_decl, params: #params_type) -> Result<#return_type> {
+                #dispatch(#ref_token ctx, &params)
             }
         },
-        // Default level — no override provided.
-        //
-        // Emit the framework-provided default impl and call its handler directly.
-        // Unlike override handlers, the default handler receives `ctx` by reference
-        // (`&` for views, `&mut` for mutable) because it is a method on the default
-        // accounts struct, not a standalone function the implementor owns.
-        (None, Some(default)) => {
-            let accounts = &default.accounts_ident;
-            let handler_fn = &default.handler_ident;
-            let ref_token = if spec.is_view { quote!(&) } else { quote!(&mut) };
-            ResolvedOverride {
-                default_impl: (default.generator)(ctx, spec),
-                fn_generics: TokenStream::new(),
-                ctx_param: ctx_param_tokens(spec.is_view, quote!(Context<#accounts>)),
-                call_with_params: quote!(#accounts::#handler_fn(#ref_token ctx, &params)),
-                call_no_params: quote!(#accounts::#handler_fn(#ref_token ctx)),
+        None => quote! {
+            pub fn #entry_fn_ident #fn_generics (#ctx_decl) -> Result<#return_type> {
+                #dispatch(#ref_token ctx)
             }
         },
-        // Unreachable: `generate_instruction` validates that at least one of
-        // override_handler or default_impl is present before calling this function.
-        (None, None) => unreachable!("validated above"),
-    }
-}
+    };
 
-/// Generates `ctx: T` or `mut ctx: T` depending on mutability.
-fn ctx_param_tokens(is_view: bool, context_type: TokenStream) -> TokenStream {
-    if is_view {
-        quote!(ctx: #context_type)
-    } else {
-        quote!(mut ctx: #context_type)
-    }
+    Ok(GeneratedInstructions { default_impls, entrypoints, extend_rules })
 }

package/src/lib.rs

@@ -5,11 +5,13 @@
 //! with optional default implementations. Each instruction is either **required** (implementor
 //! must provide a handler) or **provided** (a default implementation is available).
 //!
-//! Provided instructions support three override levels:
+//! Provided instructions support two override levels:
 //!
-//! - **Default**: No custom handler — framework emits accounts struct + entrypoint.
-//! - **Handler Override**: Custom handler reuses the default accounts struct.
-//! - **Full Override**: Custom handler with a custom accounts struct.
+//! - **Default**: No custom handler defined — the framework emits both the default accounts struct
+//!   and an entrypoint that invokes the default handler.
+//! - **Override**: Custom handler defined — the framework still emits the default accounts struct
+//!   (making it available for composition or reference), but the entrypoint delegates to the custom
+//!   handler using whatever context type it declares.
 //!
 //! See [`spec.md`](../spec.md) for the full specification.
 //!

package/src/overrides.rs

@@ -6,9 +6,9 @@
 
 use itertools::Itertools;
 use proc_macro2::TokenStream;
-use quote::quote;
+use quote::{format_ident, quote};
 use std::collections::BTreeMap;
-use syn::{spanned::Spanned, Attribute, Item, ItemFn, ItemMod, Type};
+use syn::{Attribute, Item, ItemFn, ItemMod, Type};
 
 use crate::InstructionSet;
 
@@ -25,10 +25,6 @@ pub struct OverrideConfig {
 pub struct HandlerMeta {
     /// Wrapper module identifier for the relocated handler.
     pub wrapper_mod: syn::Ident,
-    /// Original implementor-defined handler function identifier.
-    pub fn_ident: syn::Ident,
-    /// Inner type `T` extracted from `Context<...T>`.
-    pub accounts_type: Type,
     /// The full `Context<...>` type as written (preserves lifetime args).
     pub context_type: Type,
     /// Function-level generics (e.g. `<'info>`) from the override handler.
@@ -37,8 +33,11 @@ pub struct HandlerMeta {
 
 /// Collection of custom instruction handlers discovered from a module.
 pub struct InstructionOverrides<Ix: InstructionSet> {
+    /// Map of instruction variants to their handler metadata.
     handlers: BTreeMap<Ix, HandlerMeta>,
+    /// Token stream containing all generated wrapper module definitions.
     pub wrapper_defs: TokenStream,
+    /// Configuration for the override system (attribute name, domain label).
     pub config: OverrideConfig,
 }
 
@@ -52,6 +51,7 @@ impl<Ix: InstructionSet> InstructionOverrides<Ix> {
     pub fn collect(module: &ItemMod, config: OverrideConfig) -> syn::Result<Self> {
         let mut collected = Self::new(config);
 
+        // Extract the module content and early return if it's empty.
         let items = match &module.content {
             Some((_, items)) => items,
             None => return Ok(collected),
@@ -59,6 +59,7 @@ impl<Ix: InstructionSet> InstructionOverrides<Ix> {
 
         // Walk every function in the module looking for override attributes.
         for item in items {
+            // Only process functions that are annotated with the override attribute.
             let Item::Fn(fn_item) = item else { continue };
             let Some(target) = parse_override_target::<Ix>(fn_item, config)? else {
                 continue;
@@ -81,8 +82,8 @@ impl<Ix: InstructionSet> InstructionOverrides<Ix> {
             collected.handlers.insert(target, meta);
         }
 
-        // Ensure no plain `fn name()` shadows a known instruction name.
-        collected.reject_name_collisions(module)?;
+        // Ensure no unannotated `fn name()` accidentally shadows a known instruction name.
+        collected.reject_name_collisions(items)?;
         Ok(collected)
     }
 
@@ -97,10 +98,11 @@ impl<Ix: InstructionSet> InstructionOverrides<Ix> {
             return TokenStream::new();
         };
 
+        // Filter out functions that are overridden.
         let kept: Vec<_> = items
             .iter()
             .filter(|item| match item {
-                Item::Fn(fn_item) => !self.is_overridden(&fn_item.sig.ident),
+                Item::Fn(fn_item) => !has_override_attr(fn_item, self.config),
                 _ => true,
             })
             .collect();
@@ -108,26 +110,19 @@ impl<Ix: InstructionSet> InstructionOverrides<Ix> {
         quote!(#(#kept)*)
     }
 
-    /// Checks whether a function ident belongs to an override handler.
-    fn is_overridden(&self, fn_ident: &syn::Ident) -> bool {
-        self.handlers.values().any(|h| h.fn_ident == *fn_ident)
-    }
-
     /// Relocates a handler function into an isolated wrapper module
     /// and returns the corresponding [`HandlerMeta`].
     fn wrap_handler(&mut self, target: Ix, fn_item: &ItemFn) -> syn::Result<HandlerMeta> {
-        let (accounts_type, context_type) = infer_context_type(fn_item)?;
-        let fn_ident = fn_item.sig.ident.clone();
+        let context_type = infer_context_type(fn_item)?;
         let fn_generics = fn_item.sig.generics.clone();
 
         // Generate a deterministic module name: `__<attr>__<instruction>_handler`.
-        let wrapper_mod = syn::Ident::new(
-            &format!("__{}__{}_handler", self.config.attr_name, target.canonical_name()),
-            proc_macro2::Span::call_site(),
-        );
+        let wrapper_mod =
+            format_ident!("__{}__{}_handler", self.config.attr_name, target.canonical_name());
 
         // Clone the function without the instruction attribute.
-        let wrapped_fn = strip_override_attr(fn_item, self.config)?;
+        let mut wrapped_fn = fn_item.clone();
+        wrapped_fn.attrs.retain(|attr| !is_override_attr(attr, self.config));
 
         // Emit the wrapper module so the handler lives in its own scope.
         self.wrapper_defs.extend(quote! {
@@ -137,31 +132,29 @@ impl<Ix: InstructionSet> InstructionOverrides<Ix> {
             }
         });
 
-        Ok(HandlerMeta { wrapper_mod, fn_ident, accounts_type, context_type, fn_generics })
+        Ok(HandlerMeta { wrapper_mod, context_type, fn_generics })
     }
 
     /// Rejects unattributed functions whose name collides with a known instruction.
     ///
     /// A bare `fn lz_receive(...)` without `#[oapp_instruction]`
     /// would silently shadow the generated entrypoint — catch that early.
-    fn reject_name_collisions(&self, module: &ItemMod) -> syn::Result<()> {
-        let Some((_, items)) = &module.content else {
-            return Ok(());
-        };
-
+    fn reject_name_collisions(&self, items: &[Item]) -> syn::Result<()> {
         for &ix in Ix::all() {
             let name = ix.canonical_name();
-            let collides = items.iter().any(|item| match item {
-                Item::Fn(f) => f.sig.ident == name && !has_override_attr(f, self.config),
-                _ => false,
+            let collision = items.iter().find_map(|item| match item {
+                Item::Fn(f) if f.sig.ident == name && !has_override_attr(f, self.config) => {
+                    Some(&f.sig.ident)
+                },
+                _ => None,
             });
 
-            if collides {
-                return Err(syn::Error::new(
-                    module.span(),
+            if let Some(ident) = collision {
+                return Err(syn::Error::new_spanned(
+                    ident,
                     format!(
                         "`fn {name}(...)` conflicts with {domain} instruction `{name}`: \
-                         use #[{attr}] to override it",
+                        use #[{attr}] to override it",
                         domain = self.config.domain,
                         attr = self.config.attr_name,
                     ),
@@ -183,22 +176,7 @@ fn has_override_attr(fn_item: &ItemFn, cfg: OverrideConfig) -> bool {
 
 /// Returns `true` if `attr`'s last path segment matches the configured attribute name.
 fn is_override_attr(attr: &Attribute, cfg: OverrideConfig) -> bool {
-    attr.path().segments.last().is_some_and(|seg| seg.ident == cfg.attr_name)
-}
-
-/// Returns a clone of `fn_item` with the override attribute removed.
-fn strip_override_attr(fn_item: &ItemFn, cfg: OverrideConfig) -> syn::Result<ItemFn> {
-    let mut cleaned = fn_item.clone();
-    let before = cleaned.attrs.len();
-    cleaned.attrs.retain(|attr| !is_override_attr(attr, cfg));
-
-    if cleaned.attrs.len() == before {
-        return Err(syn::Error::new_spanned(
-            fn_item,
-            format!("internal error: expected #[{}] attribute on handler function", cfg.attr_name),
-        ));
-    }
-    Ok(cleaned)
+    attr.path().is_ident(cfg.attr_name)
 }
 
 /// Parses the override target from a function's override attribute.
@@ -232,11 +210,7 @@ fn parse_override_target<Ix: InstructionSet>(
     if !matches!(attr.meta, syn::Meta::Path(_)) {
         return Err(syn::Error::new_spanned(
             attr,
-            format!(
-                "#[{}] does not accept arguments; the instruction name \
-                 is inferred from the handler function name",
-                cfg.attr_name,
-            ),
+            format!("#[{}] does not accept arguments", cfg.attr_name),
         ));
     }
 
@@ -245,45 +219,46 @@ fn parse_override_target<Ix: InstructionSet>(
     let ix = Ix::from_ident(fn_ident).ok_or_else(|| {
         syn::Error::new_spanned(
             fn_ident,
-            format!("unsupported {} target `{fn_ident}`", cfg.attr_name,),
+            format!("unsupported {} target `{fn_ident}`", cfg.attr_name),
         )
     })?;
 
     Ok(Some(ix))
 }
 
-/// Extracts the inner `T` from the first `Context<T>` parameter.
+/// Extracts the `Context<...>` type from the first `&Context<T>` or `&mut Context<T>` parameter.
 ///
-/// Returns `(accounts_type, context_type)` — the inner `T` (for default-context
-/// checks) and the full `Context<...>` type as written (to preserve lifetimes in
-/// the generated entrypoint).
-fn infer_context_type(fn_item: &ItemFn) -> syn::Result<(Type, Type)> {
+/// Override handlers must accept `ctx` by reference — by-value `Context<T>` is rejected.
+/// Returns the `Context<...>` type without the reference wrapper (preserving lifetimes).
+fn infer_context_type(fn_item: &ItemFn) -> syn::Result<Type> {
     let Some(syn::FnArg::Typed(first)) = fn_item.sig.inputs.first() else {
         return Err(syn::Error::new_spanned(
             &fn_item.sig,
-            "handler must have at least one parameter of type `Context<T>`",
+            "handler must have at least one parameter of type `&Context<T>` or `&mut Context<T>`",
         ));
     };
 
-    let context_type = (*first.ty).clone();
-    let accounts_type = extract_context_inner_type(&first.ty).ok_or_else(|| {
-        syn::Error::new_spanned(&first.ty, "expected `Context<T>` as the first parameter")
-    })?;
-    Ok((accounts_type, context_type))
+    // Require `&Context<T>` or `&mut Context<T>` — reject by-value `Context<T>`.
+    let Type::Reference(ref_ty) = first.ty.as_ref() else {
+        return Err(syn::Error::new_spanned(
+            &first.ty,
+            "expected `&Context<T>` or `&mut Context<T>` as the first parameter",
+        ));
+    };
+    let inner_ty = ref_ty.elem.as_ref();
+
+    // Validate that the referenced type is `Context<...>`.
+    if !is_context_type(inner_ty) {
+        return Err(syn::Error::new_spanned(
+            &first.ty,
+            "expected `&Context<T>` or `&mut Context<T>` as the first parameter",
+        ));
+    }
+
+    Ok(inner_ty.clone())
 }
 
-/// Unwraps `Context<T>` or `Context<'a, 'b, 'c, 'info, T>` → `T`.
-///
-/// Handles both the simple `Context<T>` form and the lifetime-annotated
-/// `Context<'_, '_, '_, 'info, T>` form by finding the last type argument.
-fn extract_context_inner_type(ty: &Type) -> Option<Type> {
-    let Type::Path(tp) = ty else { return None };
-    let seg = tp.path.segments.last().filter(|s| s.ident == "Context")?;
-    let syn::PathArguments::AngleBracketed(args) = &seg.arguments else {
-        return None;
-    };
-    args.args.iter().rev().find_map(|arg| match arg {
-        syn::GenericArgument::Type(inner) => Some(inner.clone()),
-        _ => None,
-    })
+/// Returns `true` if the type's last path segment is `Context`.
+fn is_context_type(ty: &Type) -> bool {
+    matches!(ty, Type::Path(tp) if tp.path.segments.last().is_some_and(|s| s.ident == "Context"))
 }