diff --git a/docs/reference/classes/store.md b/docs/reference/classes/store.md
index aa26ee9..d8e88c8 100644
--- a/docs/reference/classes/store.md
+++ b/docs/reference/classes/store.md
@@ -13,8 +13,6 @@ Defined in: [store.ts:28](https://github.com/TanStack/store/blob/main/packages/s
 
 • **TState**
 
-• **TUpdater** *extends* `AnyUpdater` = (`cb`) => `TState`
-
 ## Constructors
 
 ### new Store()
diff --git a/docs/reference/interfaces/storeoptions.md b/docs/reference/interfaces/storeoptions.md
index 1270888..5b0e909 100644
--- a/docs/reference/interfaces/storeoptions.md
+++ b/docs/reference/interfaces/storeoptions.md
@@ -13,8 +13,6 @@ Defined in: [store.ts:5](https://github.com/TanStack/store/blob/main/packages/st
 
 • **TState**
 
-• **TUpdater** *extends* `AnyUpdater` = (`cb`) => `TState`
-
 ## Properties
 
 ### onSubscribe()?
diff --git a/examples/react/simple/src/index.tsx b/examples/react/simple/src/index.tsx
index cbf8356..2bb3cda 100644
--- a/examples/react/simple/src/index.tsx
+++ b/examples/react/simple/src/index.tsx
@@ -8,6 +8,12 @@ export const store = new Store({
   cats: 0,
 })
 
+const countStore = new Store(1, {
+    updateFn: prevState => updater => {
+        return updater(prevState) + prevState;
+    }
+});
+
 interface DisplayProps {
   animal: 'dogs' | 'cats'
 }
@@ -35,6 +41,17 @@ const Increment = ({ animal }: IncrementProps) => (
   <button onClick={() => updateState(animal)}>My Friend Likes {animal}</button>
 )
 
+const DisplayCount = () => {
+    const count = useStore(countStore);
+    return (
+        <>
+            <button onClick={() => countStore.setState(1)}>Update count - part1</button>
+            <button onClick={() => countStore.setState(() => 1)}>Update count - part2</button>
+            <div>{`count: ${count}`}</div>
+        </>
+    );
+};
+
 function App() {
   return (
     <div>
@@ -47,6 +64,7 @@ function App() {
       <Display animal="dogs" />
       <Increment animal="cats" />
       <Display animal="cats" />
+      <DisplayCount />
     </div>
   )
 }
diff --git a/packages/angular-store/src/index.ts b/packages/angular-store/src/index.ts
index af0ad0c..3fefbc8 100644
--- a/packages/angular-store/src/index.ts
+++ b/packages/angular-store/src/index.ts
@@ -17,7 +17,7 @@ export * from '@tanstack/store'
 type NoInfer<T> = [T][T extends any ? 0 : never]
 
 export function injectStore<TState, TSelected = NoInfer<TState>>(
-  store: Store<TState, any>,
+  store: Store<TState>,
   selector?: (state: NoInfer<TState>) => TSelected,
   options?: CreateSignalOptions<TSelected> & { injector?: Injector },
 ): Signal<TSelected>
@@ -27,7 +27,7 @@ export function injectStore<TState, TSelected = NoInfer<TState>>(
   options?: CreateSignalOptions<TSelected> & { injector?: Injector },
 ): Signal<TSelected>
 export function injectStore<TState, TSelected = NoInfer<TState>>(
-  store: Store<TState, any> | Derived<TState, any>,
+  store: Store<TState> | Derived<TState, any>,
   selector: (state: NoInfer<TState>) => TSelected = (d) => d as TSelected,
   options: CreateSignalOptions<TSelected> & { injector?: Injector } = {
     equal: shallow,
diff --git a/packages/react-store/src/index.ts b/packages/react-store/src/index.ts
index 155b0f5..e67d8da 100644
--- a/packages/react-store/src/index.ts
+++ b/packages/react-store/src/index.ts
@@ -9,7 +9,7 @@ export * from '@tanstack/store'
 export type NoInfer<T> = [T][T extends any ? 0 : never]
 
 export function useStore<TState, TSelected = NoInfer<TState>>(
-  store: Store<TState, any>,
+  store: Store<TState>,
   selector?: (state: NoInfer<TState>) => TSelected,
 ): TSelected
 export function useStore<TState, TSelected = NoInfer<TState>>(
@@ -17,7 +17,7 @@ export function useStore<TState, TSelected = NoInfer<TState>>(
   selector?: (state: NoInfer<TState>) => TSelected,
 ): TSelected
 export function useStore<TState, TSelected = NoInfer<TState>>(
-  store: Store<TState, any> | Derived<TState, any>,
+  store: Store<TState> | Derived<TState, any>,
   selector: (state: NoInfer<TState>) => TSelected = (d) => d as any,
 ): TSelected {
   const slice = useSyncExternalStoreWithSelector(
diff --git a/packages/solid-store/src/index.tsx b/packages/solid-store/src/index.tsx
index 3897e53..03bf45c 100644
--- a/packages/solid-store/src/index.tsx
+++ b/packages/solid-store/src/index.tsx
@@ -11,7 +11,7 @@ export * from '@tanstack/store'
 export type NoInfer<T> = [T][T extends any ? 0 : never]
 
 export function useStore<TState, TSelected = NoInfer<TState>>(
-  store: Store<TState, any>,
+  store: Store<TState>,
   selector?: (state: NoInfer<TState>) => TSelected,
 ): Accessor<TSelected>
 export function useStore<TState, TSelected = NoInfer<TState>>(
@@ -19,7 +19,7 @@ export function useStore<TState, TSelected = NoInfer<TState>>(
   selector?: (state: NoInfer<TState>) => TSelected,
 ): Accessor<TSelected>
 export function useStore<TState, TSelected = NoInfer<TState>>(
-  store: Store<TState, any> | Derived<TState, any>,
+  store: Store<TState> | Derived<TState, any>,
   selector: (state: NoInfer<TState>) => TSelected = (d) => d as any,
 ): Accessor<TSelected> {
   const [slice, setSlice] = createStore({
diff --git a/packages/store/src/store.ts b/packages/store/src/store.ts
index 94e9fc5..6cc2e26 100644
--- a/packages/store/src/store.ts
+++ b/packages/store/src/store.ts
@@ -1,15 +1,12 @@
 import { __flush } from './scheduler'
 import { isUpdaterFunction } from './types'
-import type { AnyUpdater, Listener, Updater } from './types'
+import type { Listener } from './types'
 
-export interface StoreOptions<
-  TState,
-  TUpdater extends AnyUpdater = (cb: TState) => TState,
-> {
+export interface StoreOptions< TState> {
   /**
    * Replace the default update function with a custom one.
    */
-  updateFn?: (previous: TState) => (updater: TUpdater) => TState
+  updateFn?: (previous: TState) => (updater: (prev: TState) => TState) => TState
   /**
    * Called when a listener subscribes to the store.
    *
@@ -17,7 +14,7 @@ export interface StoreOptions<
    */
   onSubscribe?: (
     listener: Listener<TState>,
-    store: Store<TState, TUpdater>,
+    store: Store<TState>,
   ) => () => void
   /**
    * Called after the state has been updated, used to derive other state.
@@ -25,16 +22,13 @@ export interface StoreOptions<
   onUpdate?: () => void
 }
 
-export class Store<
-  TState,
-  TUpdater extends AnyUpdater = (cb: TState) => TState,
-> {
+export class Store<TState> {
   listeners = new Set<Listener<TState>>()
   state: TState
   prevState: TState
-  options?: StoreOptions<TState, TUpdater>
+  options?: StoreOptions<TState>
 
-  constructor(initialState: TState, options?: StoreOptions<TState, TUpdater>) {
+  constructor(initialState: TState, options?: StoreOptions<TState>) {
     this.prevState = initialState
     this.state = initialState
     this.options = options
@@ -42,30 +36,23 @@ export class Store<
 
   subscribe = (listener: Listener<TState>) => {
     this.listeners.add(listener)
-    const unsub = this.options?.onSubscribe?.(listener, this)
+    const unsubscribe = this.options?.onSubscribe?.(listener, this)
     return () => {
       this.listeners.delete(listener)
-      unsub?.()
+      unsubscribe?.()
     }
   }
 
   /**
    * Update the store state safely with improved type checking
    */
-  setState(updater: (prevState: TState) => TState): void
-  setState(updater: TState): void
-  setState(updater: TUpdater): void
-  setState(updater: Updater<TState> | TUpdater): void {
+  setState(updater: TState | ((prevState: TState) => TState)): void {
     this.prevState = this.state
 
-    if (this.options?.updateFn) {
-      this.state = this.options.updateFn(this.prevState)(updater as TUpdater)
+    if (isUpdaterFunction(updater)) {
+      this.state = this.options?.updateFn ? this.options.updateFn(this.prevState)(updater) : updater(this.prevState)
     } else {
-      if (isUpdaterFunction(updater)) {
-        this.state = updater(this.prevState)
-      } else {
-        this.state = updater as TState
-      }
+      this.state = this.options?.updateFn ? this.options.updateFn(this.prevState)(() => updater) : updater
     }
 
     // Always run onUpdate, regardless of batching
diff --git a/packages/store/src/types.ts b/packages/store/src/types.ts
index a869cdd..be31563 100644
--- a/packages/store/src/types.ts
+++ b/packages/store/src/types.ts
@@ -1,8 +1,3 @@
-/**
- * @private
- */
-export type AnyUpdater = (prev: any) => any
-
 /**
  * Type-safe updater that can be either a function or direct value
  */
diff --git a/packages/store/tests/scheduler.test.ts b/packages/store/tests/scheduler.test.ts
index c09f564..70631b0 100644
--- a/packages/store/tests/scheduler.test.ts
+++ b/packages/store/tests/scheduler.test.ts
@@ -9,7 +9,7 @@ import {
 
 describe('Scheduler logic', () => {
   test('Should build a graph properly', () => {
-    const count = new Store<any, any>(10)
+    const count = new Store<any>(10)
 
     const halfCount = new Derived<any, any>({
       deps: [count],
@@ -36,7 +36,7 @@ describe('Scheduler logic', () => {
   })
 
   test('should unbuild a graph properly', () => {
-    const count = new Store<any, any>(10)
+    const count = new Store<any>(10)
 
     const halfCount = new Derived<any, any>({
       deps: [count],
@@ -99,7 +99,7 @@ describe('Scheduler logic', () => {
   })
 
   test('should register graph items in the wrong order properly', () => {
-    const count = new Store<any, any>(12)
+    const count = new Store<any>(12)
 
     const double = new Derived<any, any>({
       deps: [count],
@@ -124,7 +124,7 @@ describe('Scheduler logic', () => {
   })
 
   test('should register graph items in the right direction order', () => {
-    const count = new Store<any, any>(12)
+    const count = new Store<any>(12)
 
     const double = new Derived<any, any>({
       deps: [count],
diff --git a/packages/store/tests/store.test.ts b/packages/store/tests/store.test.ts
index 5d7b9ae..26db589 100644
--- a/packages/store/tests/store.test.ts
+++ b/packages/store/tests/store.test.ts
@@ -53,6 +53,24 @@ describe('store', () => {
     expect(typeof store.state).toEqual('number')
   })
 
+  test(`updateFn acts as state transformer when directly updating state`, () => {
+    const store = new Store(1, {
+      updateFn: prevState => updater => {
+        return updater(prevState) + prevState;
+      }
+    })
+
+    store.setState(1)
+
+    expect(store.state).toEqual(2)
+
+    store.setState(1)
+
+    expect(store.state).toEqual(3)
+
+    expect(typeof store.state).toEqual('number')
+  })
+
   test('listeners should receive old and new values', () => {
     const store = new Store(12)
     const fn = vi.fn()
diff --git a/packages/svelte-store/src/index.svelte.ts b/packages/svelte-store/src/index.svelte.ts
index 3d8429a..2626b7f 100644
--- a/packages/svelte-store/src/index.svelte.ts
+++ b/packages/svelte-store/src/index.svelte.ts
@@ -8,7 +8,7 @@ export * from '@tanstack/store'
 export type NoInfer<T> = [T][T extends any ? 0 : never]
 
 export function useStore<TState, TSelected = NoInfer<TState>>(
-  store: Store<TState, any>,
+  store: Store<TState>,
   selector?: (state: NoInfer<TState>) => TSelected,
 ): { readonly current: TSelected }
 export function useStore<TState, TSelected = NoInfer<TState>>(
@@ -16,7 +16,7 @@ export function useStore<TState, TSelected = NoInfer<TState>>(
   selector?: (state: NoInfer<TState>) => TSelected,
 ): { readonly current: TSelected }
 export function useStore<TState, TSelected = NoInfer<TState>>(
-  store: Store<TState, any> | Derived<TState, any>,
+  store: Store<TState> | Derived<TState, any>,
   selector: (state: NoInfer<TState>) => TSelected = (d) => d as any,
 ): { readonly current: TSelected } {
   let slice = $state(selector(store.state))
diff --git a/packages/vue-store/src/index.ts b/packages/vue-store/src/index.ts
index 0d3dfaa..ac494e2 100644
--- a/packages/vue-store/src/index.ts
+++ b/packages/vue-store/src/index.ts
@@ -10,7 +10,7 @@ export * from '@tanstack/store'
 export type NoInfer<T> = [T][T extends any ? 0 : never]
 
 export function useStore<TState, TSelected = NoInfer<TState>>(
-  store: Store<TState, any>,
+  store: Store<TState>,
   selector?: (state: NoInfer<TState>) => TSelected,
 ): Readonly<Ref<TSelected>>
 export function useStore<TState, TSelected = NoInfer<TState>>(
@@ -18,7 +18,7 @@ export function useStore<TState, TSelected = NoInfer<TState>>(
   selector?: (state: NoInfer<TState>) => TSelected,
 ): Readonly<Ref<TSelected>>
 export function useStore<TState, TSelected = NoInfer<TState>>(
-  store: Store<TState, any> | Derived<TState, any>,
+  store: Store<TState> | Derived<TState, any>,
   selector: (state: NoInfer<TState>) => TSelected = (d) => d as any,
 ): Readonly<Ref<TSelected>> {
   const slice = ref(selector(store.state)) as Ref<TSelected>