1 files changed, 13 insertions, 5 deletions

diff --git a/state/memory_state.go b/state/memory_state.go
index 95c7318..97ecad1 100644
--- a/state/memory_state.go
+++ b/state/memory_state.go
@@ -1,7 +1,7 @@
 package state
 
 import (
-"context"
+        "context"
 )
 
 type memoryStateValue struct {
@@ -13,7 +13,7 @@ func (m *memoryStateValue) Value(ctx context.Context) (any, error) {
         kg := m.backend.currentGroupKey
         key := m.backend.currentKey
         ns := m.backend.currentNamespace
-        if kgMap, exist := m.backend.states[m.stateName][kg][key]; exist {
+        if kgMap, ok := m.backend.states[m.stateName][kg][key]; ok {
                 if v, hasKey := kgMap[ns]; hasKey {
                         return v, nil
                 }
@@ -26,10 +26,15 @@ func (m *memoryStateValue) Update(ctx context.Context, newVal any) (any, error)
         kg := m.backend.currentGroupKey
         key := m.backend.currentKey
         ns := m.backend.currentNamespace
-        if _, exists := m.backend.states[m.stateName][kg][key]; !exists {
+
+        if _, ok := m.backend.states[m.stateName][kg]; !ok {
                 m.backend.states[m.stateName][kg] = make(map[string]map[any]any)
         }
 
+        if _, ok := m.backend.states[m.stateName][kg][key]; !ok {
+                m.backend.states[m.stateName][kg][key] = make(map[any]any)
+        }
+
         m.backend.states[m.stateName][kg][key][ns] = newVal
         return newVal, nil
 }
@@ -37,8 +42,11 @@ func (m *memoryStateValue) Update(ctx context.Context, newVal any) (any, error)
 func (m *memoryStateValue) Clear(ctx context.Context) error {
         kg := m.backend.currentGroupKey
         key := m.backend.currentKey
-        if kgMap, exists := m.backend.states[m.stateName][kg]; exists {
-                delete(kgMap, key)
+        ns := m.backend.currentNamespace
+        if kgMap, ok := m.backend.states[m.stateName][kg]; ok {
+                if nsMap, ok := kgMap[key]; ok {
+                        delete(nsMap, ns)
+                }
         }
         return nil
 }