diff --git a/SQL Scripts/functions/archive_context_documents_rpc.sql b/SQL Scripts/functions/archive_context_documents_rpc.sql
index 4187fb9..80b77d8 100644
--- a/SQL Scripts/functions/archive_context_documents_rpc.sql	
+++ b/SQL Scripts/functions/archive_context_documents_rpc.sql	
@@ -8,6 +8,7 @@ DECLARE
     _layer_id uuid;
     _document_id uuid;
     _row RECORD;
+    _layer_doc_id uuid;
 BEGIN
     -- Find the project for this context  
     SELECT p.id INTO _project_id FROM public.projects p 
@@ -25,22 +26,30 @@ BEGIN
     FOREACH _document_id IN ARRAY _document_ids 
     LOOP
         -- Archive the context_documents record
-        UPDATE public.context_document cd 
+        UPDATE public.context_documents cd 
           SET is_archived = TRUE 
           WHERE cd.document_id = _document_id AND cd.context_id = _context_id;
         
         -- Archive any related layers
-        FOR _row IN SELECT * FROM public.layers l 
-          INNER JOIN public.layer_contexts lc ON lc.context_id = _context_id
-          WHERE l.document_id = _document_id
+        FOR _row IN SELECT * FROM public.layer_contexts lc 
+          WHERE lc.context_id = _context_id
         LOOP 
-          UPDATE public.layers 
-            SET is_archived = TRUE 
-            WHERE id = _row.id;
+          
+          SELECT l.document_id INTO _layer_doc_id FROM public.layers l WHERE l.id = _row.layer_id;
+
+          IF _layer_doc_id = _document_id 
+          THEN
+            IF _row.is_active_layer IS TRUE 
+            THEN
+              UPDATE public.layers l
+                SET is_archived = TRUE 
+                WHERE l.id = _row.layer_id;
+            END IF;
 
-          UPDATE public.layer_contexts lc
-            SET is_archived = TRUE
-            WHERE lc.context_id = _context_id AND lc.layer_id = _row.id;
+            UPDATE public.layer_contexts lc
+              SET is_archived = TRUE
+              WHERE lc.id = _row.id;            
+          END IF;
         END LOOP;
           
     END LOOP;
diff --git a/supabase/migrations/20240822181056_fix_context_documents_edit.sql b/supabase/migrations/20240822181056_fix_context_documents_edit.sql
new file mode 100644
index 0000000..761f786
--- /dev/null
+++ b/supabase/migrations/20240822181056_fix_context_documents_edit.sql
@@ -0,0 +1,61 @@
+set check_function_bodies = off;
+
+CREATE
+OR REPLACE FUNCTION archive_context_documents_rpc (
+    _context_id uuid,
+    _document_ids uuid[] 
+) RETURNS BOOLEAN AS $body$
+DECLARE
+    _project_id uuid;
+    _layer_id uuid;
+    _document_id uuid;
+    _row RECORD;
+    _layer_doc_id uuid;
+BEGIN
+    -- Find the project for this context  
+    SELECT p.id INTO _project_id FROM public.projects p 
+      INNER JOIN public.contexts c ON c.id = _context_id 
+      WHERE p.id = c.project_id;
+
+    -- Check project policy that context documents can be updated by this user
+    IF NOT (check_action_policy_organization(auth.uid(), 'context_documents', 'UPDATE') 
+      OR check_action_policy_project(auth.uid(), 'context_documents', 'UPDATE', _project_id)) 
+    THEN
+        RETURN FALSE;
+    END IF; 
+
+    -- Iterate through the document ids and archive them in project_documents and all context_documents
+    FOREACH _document_id IN ARRAY _document_ids 
+    LOOP
+        -- Archive the context_documents record
+        UPDATE public.context_documents cd 
+          SET is_archived = TRUE 
+          WHERE cd.document_id = _document_id AND cd.context_id = _context_id;
+        
+        -- Archive any related layers
+        FOR _row IN SELECT * FROM public.layer_contexts lc 
+          WHERE lc.context_id = _context_id
+        LOOP 
+          
+          SELECT l.document_id INTO _layer_doc_id FROM public.layers l WHERE l.id = _row.layer_id;
+
+          IF _layer_doc_id = _document_id 
+          THEN
+            IF _row.is_active_layer IS TRUE 
+            THEN
+              UPDATE public.layers l
+                SET is_archived = TRUE 
+                WHERE l.id = _row.layer_id;
+            END IF;
+
+            UPDATE public.layer_contexts lc
+              SET is_archived = TRUE
+              WHERE lc.id = _row.id;            
+          END IF;
+        END LOOP;
+          
+    END LOOP;
+
+    RETURN TRUE;
+END
+$body$ LANGUAGE plpgsql SECURITY DEFINER;
\ No newline at end of file