2 files changed, 24 insertions, 15 deletions

diff --git a/state/memory_backend.go b/state/memory_backend.go
index 6594657..f4043a1 100644
--- a/state/memory_backend.go
+++ b/state/memory_backend.go
@@ -2,6 +2,7 @@ package state
 
 import (
         "context"
+
         "flux/core"
 )
 
@@ -11,9 +12,9 @@ import (
 // This should be optimize with ZERO GC Scanning, Memory Alignment
 
 type MemoryStateBackend struct {
-        maxParallelism  int
-        currentKey      string
-        currentGroupKey int
+        maxParallelism   int
+        currentKey       string
+        currentGroupKey  int
         currentNamespace any
         // 4D structure StateName -> KeyGroup -> Key -> Namespace -> Value
         states map[string]map[int]map[string]map[any]any
@@ -36,18 +37,18 @@ func (ms *MemoryStateBackend) SetCurrentKey(key string) error {
         return nil
 }
 
-func (ms *MemoryStateBackend) GetValueState(ctx context.Context, key string) (core.ValueState[any], error) {
+func (ms *MemoryStateBackend) GetValueState(ctx context.Context, statename string) (core.ValueState[any], error) {
         // Assign if state not existed
-
-        if _, ok := ms.states[key]; !ok {
-                ms.states[key] = make(map[int]map[string]map[any]any)
+        
+        if _, ok := ms.states[statename]; !ok {
+                // create state name map if not existed 
+                ms.states[statename] = make(map[int]map[string]map[any]any)
         }
 
-        return &memoryStateValue{backend: ms, stateName: key}, nil
+        return &memoryStateValue{backend: ms, stateName: statename}, nil
 }
 
-func (ms *MemoryStateBackend) GetCurrentNamespace() any{
-
+func (ms *MemoryStateBackend) GetCurrentNamespace() any {
         return ms.currentNamespace
 }
 
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
 }