diff --git a/agent-shell.el b/agent-shell.el
index 6502399d..4cdfb7e4 100644
--- a/agent-shell.el
+++ b/agent-shell.el
@@ -885,6 +885,15 @@ OUTGOING-REQUEST-DECORATOR (passed through to `acp-make-client')."
 (defvar-local agent-shell--transcript-file nil
   "Path to the shell's transcript file.")
 
+(defvar-local agent-shell--temp-directory nil
+  "When non-nil, a temporary directory to delete when the buffer is killed.")
+
+(defun agent-shell--cleanup-temp-directory ()
+  "Delete the temporary directory associated with the current buffer."
+  (when (and agent-shell--temp-directory
+             (file-directory-p agent-shell--temp-directory))
+    (delete-directory agent-shell--temp-directory t t)))
+
 (defvar agent-shell--shell-maker-config nil)
 
 ;;;###autoload
@@ -1039,10 +1048,9 @@ When NO-DISPLAY is non-nil, don't display the shell buffer."
                                                :config config
                                                :no-display no-display)))
     (with-current-buffer shell-buffer
+      (setq agent-shell--temp-directory location)
       (add-hook 'kill-buffer-hook
-                (lambda ()
-                  (when (file-directory-p location)
-                    (delete-directory location t t)))
+                #'agent-shell--cleanup-temp-directory
                 nil t))
     shell-buffer))
 
@@ -1106,12 +1114,16 @@ Works from both shell and viewport buffers."
          (config (map-elt (buffer-local-value 'agent-shell--state shell-buffer)
                           :agent-config))
          (shell-dir (buffer-local-value 'default-directory shell-buffer))
+         (temp-dir (buffer-local-value 'agent-shell--temp-directory shell-buffer))
          ;; Remember where the shell is currently displayed
          (windows (get-buffer-window-list shell-buffer nil t)))
     (with-current-buffer shell-buffer
       (when (and (agent-shell--active-requests-p (agent-shell--state))
                  (not (y-or-n-p "Agent is busy.  Restart anyway?")))
-        (user-error "Cancelled")))
+        (user-error "Cancelled"))
+      ;; Prevent the temp directory from being deleted during restart.
+      (when temp-dir
+        (remove-hook 'kill-buffer-hook #'agent-shell--cleanup-temp-directory t)))
     (kill-buffer shell-buffer)
     (let* ((default-directory shell-dir)
            (new-shell-buffer (agent-shell--start
@@ -1121,6 +1133,12 @@ Works from both shell and viewport buffers."
                               :new-session t
                               :no-focus t)))
       (shell-maker-set-buffer-name new-shell-buffer shell-buffer-name)
+      (when temp-dir
+        (with-current-buffer new-shell-buffer
+          (setq agent-shell--temp-directory temp-dir)
+          (add-hook 'kill-buffer-hook
+                    #'agent-shell--cleanup-temp-directory
+                    nil t)))
       (if (or from-viewport agent-shell-prefer-viewport-interaction)
           (agent-shell-viewport--show-buffer
            :shell-buffer new-shell-buffer)