diff --git a/module/ai/model.php b/module/ai/model.php
index f2c582d022..ea74631b74 100644
--- a/module/ai/model.php
+++ b/module/ai/model.php
@@ -1110,6 +1110,9 @@ class aiModel extends model
         }
         $linkVars = vsprintf($varsConfig->format, $vars);
 
+        /* Override method for story drafts. */
+        if($module == 'story' && $method == 'change' && !empty($object->story) && $object->story->status == 'draft') $method = 'edit';
+
         return array(helper::createLink($module, $method, $linkVars) . (empty($varsConfig->app) ? '' : "#app=$varsConfig->app"), false);
     }
 
@@ -1501,6 +1504,9 @@ class aiModel extends model
         $targetForm = $this->config->ai->targetForm[$form[0]][$form[1]];
         if(empty($targetForm)) return;
 
+        /* Override method for story drafts. */
+        if($targetForm->m == 'story' && $targetForm->f == 'change') $_SESSION['aiInjectData']['story']['edit'] = is_string($data) ? $data : json_encode($data);
+
         $_SESSION['aiInjectData'][$targetForm->m][$targetForm->f] = is_string($data) ? $data : json_encode($data);
     }
 
diff --git a/module/ai/view/inputinject.html.php b/module/ai/view/inputinject.html.php
index 0ff8b79884..d33c396cd8 100644
--- a/module/ai/view/inputinject.html.php
+++ b/module/ai/view/inputinject.html.php
@@ -25,7 +25,7 @@
   ?>
     <?php
       js::set('injectData', $injectData);
-      unset($_SESSION['aiInjectData'][$module][$method]);
+      unset($_SESSION['aiInjectData'][$module]);
       $this->app->loadLang('ai');
     ?>
     <script>
@@ -218,7 +218,7 @@
             aTag.style.display = 'none';
             document.body.appendChild(aTag);
             aTag.click();
-            
+
             if($.appCode !== 'admin') $.apps.close($.appCode);
           });
         }