Overwriting the Theme type using TypeScript (I am only using @theme-ui/css)

[Bram-Zijp]

I am trying to overwrite the Theme type as used in the @theme-ui/css module. The following example works for emotion but not for the @theme-ui/css helpers:

import type { Theme as StelaTheme } from './src';
import '@emotion/react';
import '@theme-ui/css';

declare module '@emotion/react' {
  // eslint-disable-next-line @typescript-eslint/no-empty-interface
  export interface Theme extends StelaTheme {}
}

declare module '@theme-ui/css' {
  // eslint-disable-next-line @typescript-eslint/no-empty-interface
  export interface Theme extends StelaTheme {}
}

The above example works fine for the most part.

It does not work fully for the theme inside style functions with colors for instance (and probably other properties). Colors can be undefined and it does not know about the color types I added.
Example:

boxShadow: (theme) => `0px 0px 2px ${theme.colors?.lapisBlue}`,