diff --git a/FlexLove.lua b/FlexLove.lua
index ab1f030..a842474 100644
--- a/FlexLove.lua
+++ b/FlexLove.lua
@@ -158,9 +158,9 @@ function flexlove.init(config)
     if blurOptimizations == nil then
       blurOptimizations = true -- Default to enabled
     end
-    Blur.init({ 
+    Blur.init({
       ErrorHandler = flexlove._ErrorHandler,
-      immediateModeOptimizations = blurOptimizations and config.immediateMode or false
+      immediateModeOptimizations = blurOptimizations and config.immediateMode or false,
     })
   end
 
@@ -439,11 +439,11 @@ function flexlove.endFrame()
     if element.id and element.id ~= "" then
       -- Collect state from element and all sub-modules
       local stateUpdate = element:saveState()
-      
+
       -- Use optimized update that only changes modified values
       -- Returns true if state was changed (meaning blur cache needs invalidation)
       local stateChanged = StateManager.updateStateIfChanged(element.id, stateUpdate)
-      
+
       -- Invalidate blur cache if blur-related properties changed
       if stateChanged and (element.backdropBlur or element.contentBlur) and Blur then
         Blur.clearElementCache(element.id)
diff --git a/modules/Element.lua b/modules/Element.lua
index 60b58c9..57e8030 100644
--- a/modules/Element.lua
+++ b/modules/Element.lua
@@ -465,8 +465,7 @@ function Element.new(props)
       end
     else
       -- Store as table only if non-zero values exist
-      local hasNonZero = props.cornerRadius.topLeft or props.cornerRadius.topRight or 
-                         props.cornerRadius.bottomLeft or props.cornerRadius.bottomRight
+      local hasNonZero = props.cornerRadius.topLeft or props.cornerRadius.topRight or props.cornerRadius.bottomLeft or props.cornerRadius.bottomRight
       if hasNonZero then
         self.cornerRadius = {
           topLeft = props.cornerRadius.topLeft or 0,
@@ -1458,11 +1457,6 @@ function Element.new(props)
     Element._Context.registerElement(self)
   end
 
-  -- Initialize TextEditor after element is fully constructed
-  if self._textEditor then
-    self._textEditor:restoreState(self)
-  end
-
   return self
 end
 
@@ -3172,26 +3166,15 @@ end
 ---@return ElementStateData state Complete state snapshot
 function Element:saveState()
   local state = {}
-  
-  -- Element-owned state
-  state._focused = self._focused
-  
-  -- EventHandler state (if exists)
   if self._eventHandler then
     state.eventHandler = self._eventHandler:getState()
   end
-  
-  -- TextEditor state (if exists)
   if self._textEditor then
     state.textEditor = self._textEditor:getState()
   end
-  
-  -- ScrollManager state (if exists)
   if self._scrollManager then
     state.scrollManager = self._scrollManager:getState()
   end
-  
-  -- Blur cache data (for cache invalidation)
   if self.backdropBlur or self.contentBlur then
     state.blur = {
       _blurX = self.x,
@@ -3199,18 +3182,18 @@ function Element:saveState()
       _blurWidth = self._borderBoxWidth or (self.width + self.padding.left + self.padding.right),
       _blurHeight = self._borderBoxHeight or (self.height + self.padding.top + self.padding.bottom),
     }
-    
+
     if self.backdropBlur then
       state.blur._backdropBlurIntensity = self.backdropBlur.intensity
       state.blur._backdropBlurQuality = self.backdropBlur.quality
     end
-    
+
     if self.contentBlur then
       state.blur._contentBlurIntensity = self.contentBlur.intensity
       state.blur._contentBlurQuality = self.contentBlur.quality
     end
   end
-  
+
   return state
 end
 
@@ -3221,27 +3204,28 @@ function Element:restoreState(state)
   if not state then
     return
   end
-  
-  -- Restore element-owned state
-  if state._focused ~= nil then
-    self._focused = state._focused
-  end
-  
+
   -- Restore EventHandler state (if exists)
   if self._eventHandler and state.eventHandler then
     self._eventHandler:setState(state.eventHandler)
   end
-  
+
   -- Restore TextEditor state (if exists)
   if self._textEditor and state.textEditor then
-    self._textEditor:setState(state.textEditor)
+    self._textEditor:setState(state.textEditor, self)
+    -- Sync TextEditor's focus state to Element for theme management
+    self._focused = self._textEditor._focused
+    self._cursorPosition = self._textEditor._cursorPosition
+    self._selectionStart = self._textEditor._selectionStart
+    self._selectionEnd = self._textEditor._selectionEnd
+    self._textBuffer = self._textEditor._textBuffer
   end
-  
+
   -- Restore ScrollManager state (if exists)
   if self._scrollManager and state.scrollManager then
     self._scrollManager:setState(state.scrollManager)
   end
-  
+
   -- Note: Blur cache data is used for invalidation, not restoration
 end
 
@@ -3253,10 +3237,10 @@ function Element:shouldInvalidateBlurCache(oldState, newState)
   if not oldState or not oldState.blur or not newState.blur then
     return false
   end
-  
+
   local old = oldState.blur
   local new = newState.blur
-  
+
   -- Check if any blur-related property changed
   return old._blurX ~= new._blurX
     or old._blurY ~= new._blurY
@@ -3271,31 +3255,6 @@ end
 --- Cleanup method to break circular references (for immediate mode)
 --- Note: Cleans internal module state but keeps structure for inspection
 function Element:_cleanup()
-  -- Clean up module internal state
-  if self._eventHandler then
-    self._eventHandler:_cleanup()
-  end
-  
-  if self._themeManager then
-    self._themeManager:_cleanup()
-  end
-  
-  if self._renderer then
-    self._renderer:_cleanup()
-  end
-  
-  if self._layoutEngine then
-    self._layoutEngine:_cleanup()
-  end
-  
-  if self._scrollManager then
-    self._scrollManager:_cleanup()
-  end
-  
-  if self._textEditor then
-    self._textEditor:_cleanup()
-  end
-  
   -- Clear event callbacks (may hold closures)
   self.onEvent = nil
   self.onFocus = nil
diff --git a/modules/EventHandler.lua b/modules/EventHandler.lua
index 31b590a..2108099 100644
--- a/modules/EventHandler.lua
+++ b/modules/EventHandler.lua
@@ -218,7 +218,6 @@ end
 ---@param my number Mouse Y position
 ---@param button number Mouse button (1=left, 2=right, 3=middle)
 function EventHandler:_handleMousePress(element, mx, my, button)
-
   -- Check if press is on scrollbar first (skip if already handled)
   if button == 1 and not self._scrollbarPressHandled and element._handleScrollbarPress then
     if element:_handleScrollbarPress(mx, my, button) then
@@ -263,7 +262,6 @@ end
 ---@param button number Mouse button
 ---@param isHovering boolean Whether mouse is over element
 function EventHandler:_handleMouseDrag(element, mx, my, button, isHovering)
-
   local lastX = self._lastMouseX[button] or mx
   local lastY = self._lastMouseY[button] or my
 
@@ -303,7 +301,6 @@ end
 ---@param my number Mouse Y position
 ---@param button number Mouse button
 function EventHandler:_handleMouseRelease(element, mx, my, button)
-
   local currentTime = love.timer.getTime()
   local modifiers = EventHandler._utils.getModifiers()
 
@@ -471,7 +468,6 @@ end
 ---@param y number Touch Y position
 ---@param pressure number Touch pressure (0-1)
 function EventHandler:_handleTouchBegan(element, touchId, x, y, pressure)
-
   -- Create touch state
   self._touches[touchId] = {
     x = x,
@@ -642,13 +638,4 @@ function EventHandler:_invokeCallback(element, event)
   end
 end
 
-
---- Cleanup method to break circular references (for immediate mode)
---- Note: Only clears module references, preserves state for inspection/testing
-function EventHandler:_cleanup()
-  -- DO NOT clear state data (_pressed, _touches, etc.) - they're needed for state persistence
-  -- Only clear module references that could create circular dependencies
-  -- (In practice, EventHandler doesn't store refs to Context/utils, so nothing to do)
-end
-
 return EventHandler
diff --git a/modules/LayoutEngine.lua b/modules/LayoutEngine.lua
index 2c2c352..0348c1f 100644
--- a/modules/LayoutEngine.lua
+++ b/modules/LayoutEngine.lua
@@ -172,7 +172,7 @@ function LayoutEngine:layoutChildren()
     timerName = "layout_" .. (self.element.id or tostring(self.element):match("0x%x+") or "unknown")
     LayoutEngine._Performance:startTimer(timerName)
   end
-  
+
   if self.element == nil then
     return
   end
@@ -1023,12 +1023,14 @@ function LayoutEngine:_canSkipLayout()
   end
 
   local cache = self._layoutCache
-  
+
   -- Check if layout inputs have changed
-  if cache.childrenCount == childrenCount and
-     cache.containerWidth == containerWidth and
-     cache.containerHeight == containerHeight and
-     cache.childrenHash == childrenHash then
+  if
+    cache.childrenCount == childrenCount
+    and cache.containerWidth == containerWidth
+    and cache.containerHeight == containerHeight
+    and cache.childrenHash == childrenHash
+  then
     return true -- Layout hasn't changed, can skip
   end
 
@@ -1072,13 +1074,4 @@ function LayoutEngine:_trackLayoutRecalculation()
   end
 end
 
-
---- Cleanup method to break circular references (for immediate mode)
-function LayoutEngine:_cleanup()
-  -- Circular refs: Element → LayoutEngine → element → Element
-  -- But breaking element ref breaks functionality
-  -- Module refs are singletons, not circular
-  -- In immediate mode, full GC happens anyway
-end
-
 return LayoutEngine
diff --git a/modules/Renderer.lua b/modules/Renderer.lua
index 4d91ccc..eb4db99 100644
--- a/modules/Renderer.lua
+++ b/modules/Renderer.lua
@@ -120,8 +120,6 @@ function Renderer.new(config, deps)
   return self
 end
 
-
-
 --- Get or create blur instance for this element
 ---@return table|nil Blur instance or nil
 function Renderer:getBlurInstance()
@@ -202,10 +200,7 @@ function Renderer:_drawImage(x, y, paddingLeft, paddingTop, contentWidth, conten
     if type(self.cornerRadius) == "number" then
       hasCornerRadius = self.cornerRadius > 0
     else
-      hasCornerRadius = self.cornerRadius.topLeft > 0
-        or self.cornerRadius.topRight > 0
-        or self.cornerRadius.bottomLeft > 0
-        or self.cornerRadius.bottomRight > 0
+      hasCornerRadius = self.cornerRadius.topLeft > 0 or self.cornerRadius.topRight > 0 or self.cornerRadius.bottomLeft > 0 or self.cornerRadius.bottomRight > 0
     end
   end
 
@@ -340,7 +335,7 @@ function Renderer:_drawBorders(x, y, borderBoxWidth, borderBoxHeight)
   if not self.border then
     return
   end
-  
+
   -- Handle border as number (uniform border width)
   if type(self.border) == "number" then
     local borderColorWithOpacity = self._Color.new(self.borderColor.r, self.borderColor.g, self.borderColor.b, self.borderColor.a * self.opacity)
@@ -348,7 +343,7 @@ function Renderer:_drawBorders(x, y, borderBoxWidth, borderBoxHeight)
     self._RoundedRect.draw("line", x, y, borderBoxWidth, borderBoxHeight, self.cornerRadius)
     return
   end
-  
+
   local borderColorWithOpacity = self._Color.new(self.borderColor.r, self.borderColor.g, self.borderColor.b, self.borderColor.a * self.opacity)
   love.graphics.setColor(borderColorWithOpacity:toRGBA())
 
@@ -963,12 +958,4 @@ function Renderer:destroy()
   self._blurInstance = nil
 end
 
-
---- Cleanup method to break circular references (for immediate mode)
-function Renderer:_cleanup()
-  -- Renderer doesn't create circular references (no element back-reference)
-  -- Module refs are singletons
-  -- In immediate mode, full element cleanup happens via array clearing
-end
-
 return Renderer
diff --git a/modules/ScrollManager.lua b/modules/ScrollManager.lua
index ee6710c..4b63d1b 100644
--- a/modules/ScrollManager.lua
+++ b/modules/ScrollManager.lua
@@ -640,35 +640,35 @@ function ScrollManager:setState(state)
   elseif state.scrollX ~= nil then
     self._scrollX = state.scrollX
   end
-  
+
   if state._scrollY ~= nil then
     self._scrollY = state._scrollY
   elseif state.scrollY ~= nil then
     self._scrollY = state.scrollY
   end
-  
+
   if state._scrollbarDragging ~= nil then
     self._scrollbarDragging = state._scrollbarDragging
   elseif state.scrollbarDragging ~= nil then
     self._scrollbarDragging = state.scrollbarDragging
   end
-  
+
   if state._hoveredScrollbar ~= nil then
     self._hoveredScrollbar = state._hoveredScrollbar
   elseif state.hoveredScrollbar ~= nil then
     self._hoveredScrollbar = state.hoveredScrollbar
   end
-  
+
   if state._scrollbarDragOffset ~= nil then
     self._scrollbarDragOffset = state._scrollbarDragOffset
   elseif state.scrollbarDragOffset ~= nil then
     self._scrollbarDragOffset = state.scrollbarDragOffset
   end
-  
+
   if state._scrollbarHoveredVertical ~= nil then
     self._scrollbarHoveredVertical = state._scrollbarHoveredVertical
   end
-  
+
   if state._scrollbarHoveredHorizontal ~= nil then
     self._scrollbarHoveredHorizontal = state._scrollbarHoveredHorizontal
   end
@@ -894,14 +894,4 @@ function ScrollManager:isMomentumScrolling()
   return self._momentumScrolling
 end
 
-
---- Cleanup method to break circular references (for immediate mode)
-function ScrollManager:_cleanup()
-  -- Cleanup breaks circular references only
-  -- The main circular ref is: Element → ScrollManager → element → Element
-  -- Breaking element ref would break functionality, so we keep it
-  -- Module refs (_utils, _Color) are not circular, they're shared singletons
-  -- In immediate mode, the whole element will be GC'd anyway, so minimal cleanup needed
-end
-
 return ScrollManager
diff --git a/modules/TextEditor.lua b/modules/TextEditor.lua
index 6ce9093..9c47134 100644
--- a/modules/TextEditor.lua
+++ b/modules/TextEditor.lua
@@ -1123,7 +1123,6 @@ function TextEditor:handleTextInput(element, text)
 
   -- Insert text at cursor position
   self:insertText(element, text)
-
   -- Trigger onTextChange callback
   if self.onTextChange and self._textBuffer ~= oldText then
     self.onTextChange(element, self._textBuffer, oldText)
@@ -1699,45 +1698,50 @@ end
 
 --- Restore state from persistence
 ---@param state table State to restore
-function TextEditor:setState(state)
+---@param element Element? The parent element (needed for focus restoration)
+function TextEditor:setState(state, element)
   if not state then
     return
   end
-  
+
   if state._cursorPosition ~= nil then
     self._cursorPosition = state._cursorPosition
   end
-  
+
   if state._selectionStart ~= nil then
     self._selectionStart = state._selectionStart
   end
-  
+
   if state._selectionEnd ~= nil then
     self._selectionEnd = state._selectionEnd
   end
-  
+
   if state._textBuffer ~= nil then
     self._textBuffer = state._textBuffer
   end
-  
+
   if state._cursorBlinkTimer ~= nil then
     self._cursorBlinkTimer = state._cursorBlinkTimer
   end
-  
+
   if state._cursorVisible ~= nil then
     self._cursorVisible = state._cursorVisible
   end
-  
+
   if state._cursorBlinkPaused ~= nil then
     self._cursorBlinkPaused = state._cursorBlinkPaused
   end
-  
+
   if state._cursorBlinkPauseTimer ~= nil then
     self._cursorBlinkPauseTimer = state._cursorBlinkPauseTimer
   end
-  
+
   if state._focused ~= nil then
     self._focused = state._focused
+    -- Restore focused element in Context if this element was focused
+    if self._focused and element then
+      self._Context._focusedElement = element
+    end
   end
 end
 
@@ -1748,7 +1752,12 @@ function TextEditor:_saveState(element)
     return
   end
 
-  self._StateManager.updateState(element._stateId, {
+  -- Get current state (may have other sub-modules like eventHandler, scrollManager)
+  local currentState = self._StateManager.getState(element._stateId) or {}
+
+  -- Update only the textEditor sub-table to match the nested structure
+  -- used by element:saveState() at endFrame
+  currentState.textEditor = {
     _focused = self._focused,
     _textBuffer = self._textBuffer,
     _cursorPosition = self._cursorPosition,
@@ -1758,14 +1767,9 @@ function TextEditor:_saveState(element)
     _cursorVisible = self._cursorVisible,
     _cursorBlinkPaused = self._cursorBlinkPaused,
     _cursorBlinkPauseTimer = self._cursorBlinkPauseTimer,
-  })
-end
+  }
 
-
---- Cleanup method to break circular references (for immediate mode)
-function TextEditor:_cleanup()
-  -- TextEditor → element is circular, but breaking it breaks functionality
-  -- Module refs are singletons, not circular
+  self._StateManager.updateState(element._stateId, currentState)
 end
 
 return TextEditor
diff --git a/modules/Theme.lua b/modules/Theme.lua
index 8ad6e32..1a3139d 100644
--- a/modules/Theme.lua
+++ b/modules/Theme.lua
@@ -961,14 +961,6 @@ function ThemeManager:setTheme(themeName, componentName)
   self.themeComponent = componentName
 end
 
-
---- Cleanup method to break circular references (for immediate mode)
-function ThemeManager:_cleanup()
-  -- ThemeManager doesn't create circular references
-  -- Theme refs are to shared theme objects
-end
-
--- Export both Theme and ThemeManager
 Theme.Manager = ThemeManager
 
 --- Check theme definitions for correctness before use to catch configuration errors early
diff --git a/testing/__tests__/text_editor_test.lua b/testing/__tests__/text_editor_test.lua
index b703867..5cd53b9 100644
--- a/testing/__tests__/text_editor_test.lua
+++ b/testing/__tests__/text_editor_test.lua
@@ -1940,7 +1940,8 @@ function TestTextEditorStateSaving:test_saveState_immediate_mode()
   editor:setText(element, "New text")
 
   luaunit.assertNotNil(savedState)
-  luaunit.assertEquals(savedState._textBuffer, "New text")
+  luaunit.assertNotNil(savedState.textEditor, "State should have textEditor sub-table")
+  luaunit.assertEquals(savedState.textEditor._textBuffer, "New text")
 end
 
 function TestTextEditorStateSaving:test_saveState_not_immediate_mode()
@@ -1973,6 +1974,95 @@ function TestTextEditorStateSaving:test_saveState_not_immediate_mode()
   luaunit.assertFalse(saveCalled)
 end
 
+-- ============================================================================
+-- Immediate Mode State Persistence Tests
+-- ============================================================================
+
+TestTextEditorImmediateMode = {}
+
+function TestTextEditorImmediateMode:test_focus_persists_across_frames()
+  -- Simulate immediate mode
+  MockContext._immediateMode = true
+  
+  -- Create editor and focus it
+  local editor1 = createTextEditor({text = "Hello", editable = true})
+  local element1 = createMockElement()
+  editor1:focus(element1)
+  
+  luaunit.assertTrue(editor1._focused)
+  luaunit.assertEquals(MockContext._focusedElement, element1)
+  
+  -- Get state
+  local state = editor1:getState()
+  luaunit.assertTrue(state._focused)
+  
+  -- Create new editor (simulating next frame)
+  local editor2 = createTextEditor({text = "Hello", editable = true})
+  local element2 = createMockElement()
+  
+  -- Restore state
+  editor2:setState(state, element2)
+  
+  -- Focus should be restored
+  luaunit.assertTrue(editor2._focused)
+  luaunit.assertEquals(MockContext._focusedElement, element2)
+  
+  -- Reset
+  MockContext._immediateMode = false
+  MockContext._focusedElement = nil
+end
+
+function TestTextEditorImmediateMode:test_text_input_works_after_state_restore()
+  -- Simulate immediate mode
+  MockContext._immediateMode = true
+  
+  -- Create editor and focus it
+  local editor1 = createTextEditor({text = "Hello", editable = true})
+  local element1 = createMockElement()
+  editor1:focus(element1)
+  
+  -- Get state
+  local state = editor1:getState()
+  
+  -- Create new editor (simulating next frame)
+  local editor2 = createTextEditor({text = "Hello", editable = true})
+  local element2 = createMockElement()
+  
+  -- Restore state (this should restore focus)
+  editor2:setState(state, element2)
+  
+  -- Text input should work because focus was restored
+  editor2:handleTextInput(element2, "X")
+  luaunit.assertEquals(editor2:getText(), "HelloX")
+  
+  -- Reset
+  MockContext._immediateMode = false
+  MockContext._focusedElement = nil
+end
+
+function TestTextEditorImmediateMode:test_cursor_position_persists()
+  MockContext._immediateMode = true
+  
+  local editor1 = createTextEditor({text = "Hello", editable = true})
+  local element1 = createMockElement()
+  editor1:focus(element1)
+  editor1:moveCursorBy(element1, -2) -- Move cursor to position 3
+  
+  local state = editor1:getState()
+  luaunit.assertEquals(state._cursorPosition, 3)
+  
+  -- Restore in new editor
+  local editor2 = createTextEditor({text = "Hello", editable = true})
+  local element2 = createMockElement()
+  editor2:setState(state, element2)
+  
+  luaunit.assertEquals(editor2._cursorPosition, 3)
+  
+  -- Reset
+  MockContext._immediateMode = false
+  MockContext._focusedElement = nil
+end
+
 -- ============================================================================
 -- Run Tests
 -- ============================================================================