diff --git a/caido.config.ts b/caido.config.ts
index 0e2d053..15e1c77 100644
--- a/caido.config.ts
+++ b/caido.config.ts
@@ -16,7 +16,7 @@ export default defineConfig({
   id,
   name: "Shift",
   description: "Delegate your work to Shift",
-  version: "2.1.0",
+  version: "2.1.1",
   author: {
     name: "Caido Labs Inc.",
     email: "dev@caido.io",
diff --git a/packages/frontend/src/agents/types/index.ts b/packages/frontend/src/agents/types/index.ts
index ed0557e..da16d36 100644
--- a/packages/frontend/src/agents/types/index.ts
+++ b/packages/frontend/src/agents/types/index.ts
@@ -1,4 +1,4 @@
-export * from "./config";
 export * from "./agent";
-export * from "./todos";
+export * from "./config";
 export * from "./runtime";
+export * from "./todos";
diff --git a/packages/frontend/src/components/float/Container.vue b/packages/frontend/src/components/float/Container.vue
index b8e37b1..a4a8298 100644
--- a/packages/frontend/src/components/float/Container.vue
+++ b/packages/frontend/src/components/float/Container.vue
@@ -1,5 +1,7 @@
 <script setup lang="ts">
+import { onKeyStroke } from "@vueuse/core";
 import { storeToRefs } from "pinia";
+import { onMounted, ref } from "vue";
 
 import Actions from "@/components/float/Actions.vue";
 import { useDragResize } from "@/components/float/useDragResize";
@@ -17,6 +19,17 @@ const { style, onDragMouseDown } = useDragResize({
 
 const store = useFloatStore();
 const { isRunning, query } = storeToRefs(store);
+
+const textarea = ref<HTMLTextAreaElement>();
+
+onMounted(() => {
+  store.textarea = textarea.value;
+  textarea.value?.focus();
+});
+
+onKeyStroke("Escape", () => {
+  store.closeFloat();
+});
 </script>
 
 <template>
diff --git a/packages/frontend/src/components/float/useDragResize.ts b/packages/frontend/src/components/float/useDragResize.ts
index 58e0283..4c35c7b 100644
--- a/packages/frontend/src/components/float/useDragResize.ts
+++ b/packages/frontend/src/components/float/useDragResize.ts
@@ -1,23 +1,19 @@
 import { useEventListener } from "@vueuse/core";
 import { computed, ref } from "vue";
 
+const WIDTH = 500;
+const HEIGHT = 125;
+
+const clamp = (value: number, min: number, max: number) =>
+  Math.max(min, Math.min(value, max));
+
 export const useDragResize = (options: {
   initialTop: number;
   initialLeft: number;
 }) => {
-  const width = ref(500);
-  const height = ref(125);
   const top = ref(options.initialTop);
   const left = ref(options.initialLeft);
 
-  let stopMove: (() => void) | undefined = undefined;
-  let stopUp: (() => void) | undefined = undefined;
-  const isResizing = ref(false);
-  const startX = ref(0);
-  const startY = ref(0);
-  const startWidth = ref(0);
-  const startHeight = ref(0);
-
   let stopMoveDrag: (() => void) | undefined = undefined;
   let stopUpDrag: (() => void) | undefined = undefined;
   const isDragging = ref(false);
@@ -26,52 +22,18 @@ export const useDragResize = (options: {
   const startTop = ref(0);
   const startLeft = ref(0);
 
-  const onMouseMove = (event: MouseEvent) => {
-    if (!isResizing.value) {
-      return;
-    }
-    const deltaX = event.clientX - startX.value;
-    const deltaY = event.clientY - startY.value;
-    const nextWidth = startWidth.value + deltaX;
-    const nextHeight = startHeight.value + deltaY;
-    width.value = nextWidth > 300 ? nextWidth : 300;
-    height.value = nextHeight > 120 ? nextHeight : 120;
-  };
-
-  const onMouseUp = () => {
-    if (!isResizing.value) {
-      return;
-    }
-    isResizing.value = false;
-    if (stopMove !== undefined) {
-      stopMove();
-      stopMove = undefined;
-    }
-    if (stopUp !== undefined) {
-      stopUp();
-      stopUp = undefined;
-    }
-  };
-
-  const onResizeMouseDown = (event: MouseEvent) => {
-    isResizing.value = true;
-    startX.value = event.clientX;
-    startY.value = event.clientY;
-    startWidth.value = width.value;
-    startHeight.value = height.value;
-    stopMove = useEventListener(document, "mousemove", onMouseMove);
-    stopUp = useEventListener(document, "mouseup", onMouseUp);
-    event.preventDefault();
-  };
-
   const onDragMouseMove = (event: MouseEvent) => {
     if (!isDragging.value) {
       return;
     }
     const deltaX = event.clientX - dragStartX.value;
     const deltaY = event.clientY - dragStartY.value;
-    left.value = startLeft.value + deltaX;
-    top.value = startTop.value + deltaY;
+
+    const maxLeft = window.innerWidth - WIDTH;
+    const maxTop = window.innerHeight - HEIGHT;
+
+    left.value = clamp(startLeft.value + deltaX, 0, maxLeft);
+    top.value = clamp(startTop.value + deltaY, 0, maxTop);
   };
 
   const onDragMouseUp = () => {
@@ -90,9 +52,6 @@ export const useDragResize = (options: {
   };
 
   const onDragMouseDown = (event: MouseEvent) => {
-    if (isResizing.value) {
-      return;
-    }
     isDragging.value = true;
     dragStartX.value = event.clientX;
     dragStartY.value = event.clientY;
@@ -103,18 +62,15 @@ export const useDragResize = (options: {
     event.preventDefault();
   };
 
-  const style = computed(() => {
-    return {
-      top: `${top.value}px`,
-      left: `${left.value}px`,
-      width: `${width.value}px`,
-      height: `${height.value}px`,
-    };
-  });
+  const style = computed(() => ({
+    top: `${top.value}px`,
+    left: `${left.value}px`,
+    width: `${WIDTH}px`,
+    height: `${HEIGHT}px`,
+  }));
 
   return {
     style,
-    onResizeMouseDown,
     onDragMouseDown,
   };
 };
diff --git a/packages/frontend/src/float/float.ts b/packages/frontend/src/float/float.ts
index e305ab0..c5bf6c9 100644
--- a/packages/frontend/src/float/float.ts
+++ b/packages/frontend/src/float/float.ts
@@ -25,6 +25,9 @@ export const setupFloat = (sdk: FrontendSDK) => {
   });
 };
 
+const FLOAT_WIDTH = 500;
+const FLOAT_HEIGHT = 125;
+
 const spawnFloat = (sdk: FrontendSDK) => {
   if (document.querySelector("[data-plugin='shift-float']")) {
     return;
@@ -37,8 +40,12 @@ const spawnFloat = (sdk: FrontendSDK) => {
   container.style.zIndex = "3000";
   document.body.appendChild(container);
 
-  const initialLeft = lastCursorX !== undefined ? lastCursorX - 250 : 0;
-  const initialTop = lastCursorY !== undefined ? lastCursorY - 50 : 0;
+  const maxLeft = window.innerWidth - FLOAT_WIDTH;
+  const maxTop = window.innerHeight - FLOAT_HEIGHT;
+  const rawLeft = lastCursorX !== undefined ? lastCursorX - 250 : 0;
+  const rawTop = lastCursorY !== undefined ? lastCursorY - 50 : 0;
+  const initialLeft = Math.max(0, Math.min(rawLeft, maxLeft));
+  const initialTop = Math.max(0, Math.min(rawTop, maxTop));
 
   const app = createApp(ShiftFloat, {
     initialTop,
diff --git a/packages/frontend/src/stores/float.ts b/packages/frontend/src/stores/float.ts
index f4ac056..e1a57d6 100644
--- a/packages/frontend/src/stores/float.ts
+++ b/packages/frontend/src/stores/float.ts
@@ -32,13 +32,17 @@ export const useFloatStore = defineStore("stores.float", () => {
   };
 
   const handleKeydown = (event: KeyboardEvent) => {
-    const isTextareaFocused = document.activeElement === textarea.value;
-
     if (event.key === "Escape") {
       closeFloat();
     }
 
-    if (event.key === "ArrowUp" && isTextareaFocused) {
+    if (event.key === "Enter" && !event.shiftKey) {
+      runQuery();
+      event.preventDefault();
+      return;
+    }
+
+    if (event.key === "ArrowUp") {
       const textareaElement = textarea.value;
       if (textareaElement !== undefined) {
         const caretIndex = textareaElement.selectionStart ?? 0;
@@ -66,17 +70,11 @@ export const useFloatStore = defineStore("stores.float", () => {
           }
         }
       }
-
       event.stopPropagation();
       return;
     }
 
-    if (event.key === "Enter" && !event.shiftKey && isTextareaFocused) {
-      runQuery();
-      event.preventDefault();
-    }
-
-    if (event.key === "ArrowUp" || event.key === "ArrowDown") {
+    if (event.key === "ArrowDown") {
       event.stopPropagation();
       return;
     }