diff --git a/.size-snapshot.json b/.size-snapshot.json
index 1dbf62fe3c..7510c83b64 100644
--- a/.size-snapshot.json
+++ b/.size-snapshot.json
@@ -68,9 +68,9 @@
     "gzipped": 251
   },
   "middleware.js": {
-    "bundled": 2039,
-    "minified": 1207,
-    "gzipped": 610
+    "bundled": 2088,
+    "minified": 1231,
+    "gzipped": 628
   },
   "vanilla.js": {
     "bundled": 1990,
diff --git a/src/middleware.ts b/src/middleware.ts
index d4efa5eadf..20d6b8803a 100644
--- a/src/middleware.ts
+++ b/src/middleware.ts
@@ -1,21 +1,42 @@
-const redux = (reducer: any, initial: any) => (
-  set: any,
-  get: any,
-  api: any
-) => {
-  api.dispatch = (action: any) => {
-    set((state: any) => reducer(state, action))
+import type {
+  State,
+  GetState,
+  SetState,
+  StoreApi,
+  PartialState,
+} from './vanilla'
+
+const redux = <S extends State, A extends { type: unknown }>(
+  reducer: (state: S, action: A) => S,
+  initial: S
+) => (
+  set: SetState<S & { dispatch: (a: A) => A }>,
+  get: GetState<S & { dispatch: (a: A) => A }>,
+  api: StoreApi<S & { dispatch: (a: A) => A }> & {
+    dispatch?: (a: A) => A
+    devtools?: any
+  }
+): S & { dispatch: (a: A) => A } => {
+  api.dispatch = (action: A) => {
+    set((state: S) => reducer(state, action))
     api.devtools && api.devtools.send(api.devtools.prefix + action.type, get())
     return action
   }
   return { dispatch: api.dispatch, ...initial }
 }
 
-const devtools = (fn: any, prefix?: string) => (
-  set: any,
-  get: any,
-  api: any
-) => {
+const devtools = <S extends State>(
+  fn: (
+    set: (partial: PartialState<S>, replace?: boolean, name?: string) => void,
+    get: GetState<S>,
+    api: StoreApi<S>
+  ) => S,
+  prefix?: string
+) => (
+  set: SetState<S>,
+  get: GetState<S>,
+  api: StoreApi<S> & { devtools?: any }
+): S => {
   let extension
   try {
     extension =
@@ -30,10 +51,14 @@ const devtools = (fn: any, prefix?: string) => (
     api.devtools = null
     return fn(set, get, api)
   }
-  const namedSet = (state: any, name?: any) => {
-    set(state)
-    if (name) {
-      api.devtools.send(api.devtools.prefix + name, get())
+  const namedSet = (
+    state: PartialState<S>,
+    replace?: boolean,
+    name?: string | false
+  ) => {
+    set(state, replace)
+    if (name !== false) {
+      api.devtools.send(api.devtools.prefix + (name || 'action'), get())
     }
   }
   const initialState = fn(namedSet, get, api)
@@ -47,6 +72,7 @@ const devtools = (fn: any, prefix?: string) => (
           message.payload.type === 'JUMP_TO_STATE'
         namedSet(
           JSON.parse(message.state),
+          false,
           !initialState.dispatch && !ignoreState && 'setState'
         )
       }