How do I register a custom extension for offchain workers?

I have the following code.

#[cfg(feature = "std")]
use sp_externalities::ExternalitiesExt;

#[cfg(feature = "std")]
sp_externalities::decl_extension! {
    pub struct CustomExt(u32);
}

#[cfg(feature = "std")]
pub type HostFunctions = (custom::HostFunctions,);

#[runtime_interface]
pub trait Custom {
    fn get_val(&mut self) -> Option<u32> {
        self.extension::<CustomExt>().map(|ext| ext.0)
    }
}

I've also passed the extension to custom_extension in service.rs as below

if enable_offchain_worker {
    task_manager.spawn_handle().spawn(
        "offchain-workers-runner",
        "offchain-work",
        sc_offchain::OffchainWorkers::new(sc_offchain::OffchainWorkerOptions {
            ...
            custom_extensions: move |_| {
                vec![
                    Box::new(statement_store.clone().as_statement_store_ext()) as Box<_>,
                    // This should register my extension.
                    Box::new(CustomExt(10)) as Box<_>,
                ]
            },
        })
        .run(client.clone(), task_manager.spawn_handle())
        .boxed(),
    );
}

And I've added the host functions in executor.rs

impl sc_executor::NativeExecutionDispatch for ExecutorDispatch {
    type ExtendHostFunctions = (
        frame_benchmarking::benchmarking::HostFunctions,
        sp_statement_store::runtime_api::HostFunctions,
        my_custom_ext::HostFunctions,
    );

    fn dispatch(method: &str, data: &[u8]) -> Option<Vec<u8>> {
        // Omitted
    }

    fn native_version() -> sc_executor::NativeVersion {
        // Omitted
    }
}

But when I try to access the value, I always get None

if let Some(key) = my_custom_ext::get_val() {
    log::info!("Got value {key}");
} else {
    log::warn!("Got None");
}

I can't seem to figure out why the extension is not working, is there anything else that I am missing?

Thanks

It turned out to be a bug in substrate. It is already fixed in this PR.