Fix context key matching

* You need to check all layers of the context stack
* When in command, the context should be based on where focus was (to
  match `available_actions`.

crates/command_palette2/src/command_palette.rs

@@ -311,7 +311,11 @@ impl PickerDelegate for CommandPaletteDelegate {
                         command.name.clone(),
                         r#match.positions.clone(),
                     ))
-                    .children(KeyBinding::for_action(&*command.action, cx)),
+                    .children(KeyBinding::for_action_in(
+                        &*command.action,
+                        &self.previous_focus_handle,
+                        cx,
+                    )),
             ),
         )
     }

crates/gpui2/src/key_dispatch.rs

@@ -16,7 +16,7 @@ pub struct DispatchNodeId(usize);
 
 pub(crate) struct DispatchTree {
     node_stack: Vec<DispatchNodeId>,
-    context_stack: Vec<KeyContext>,
+    pub(crate) context_stack: Vec<KeyContext>,
     nodes: Vec<DispatchNode>,
     focusable_node_ids: HashMap<FocusId, DispatchNodeId>,
     keystroke_matchers: HashMap<SmallVec<[KeyContext; 4]>, KeystrokeMatcher>,
@@ -163,13 +163,24 @@ impl DispatchTree {
         actions
     }
 
-    pub fn bindings_for_action(&self, action: &dyn Action) -> Vec<KeyBinding> {
+    pub fn bindings_for_action(
+        &self,
+        action: &dyn Action,
+        context_stack: &Vec<KeyContext>,
+    ) -> Vec<KeyBinding> {
         self.keymap
             .lock()
             .bindings_for_action(action.type_id())
             .filter(|candidate| {
-                candidate.action.partial_eq(action)
+                if !candidate.action.partial_eq(action) {
-                    && candidate.matches_context(&self.context_stack)
+                    return false;
+                }
+                for i in 1..context_stack.len() {
+                    if candidate.matches_context(&context_stack[0..i]) {
+                        return true;
+                    }
+                }
+                return false;
             })
             .cloned()
             .collect()

crates/gpui2/src/window.rs

@@ -1492,10 +1492,28 @@ impl<'a> WindowContext<'a> {
     }
 
     pub fn bindings_for_action(&self, action: &dyn Action) -> Vec<KeyBinding> {
-        self.window
+        self.window.current_frame.dispatch_tree.bindings_for_action(
-            .current_frame
+            action,
-            .dispatch_tree
+            &self.window.current_frame.dispatch_tree.context_stack,
-            .bindings_for_action(action)
+        )
+    }
+
+    pub fn bindings_for_action_in(
+        &self,
+        action: &dyn Action,
+        focus_handle: &FocusHandle,
+    ) -> Vec<KeyBinding> {
+        let dispatch_tree = &self.window.previous_frame.dispatch_tree;
+
+        let Some(node_id) = dispatch_tree.focusable_node_id(focus_handle.id) else {
+            return vec![];
+        };
+        let context_stack = dispatch_tree
+            .dispatch_path(node_id)
+            .into_iter()
+            .map(|node_id| dispatch_tree.node(node_id).context.clone())
+            .collect();
+        dispatch_tree.bindings_for_action(action, &context_stack)
     }
 
     pub fn listener_for<V: Render, E>(

crates/ui2/src/components/keybinding.rs

@@ -1,5 +1,5 @@
 use crate::{h_stack, prelude::*, Icon, IconElement, IconSize};
-use gpui::{relative, rems, Action, Div, IntoElement, Keystroke};
+use gpui::{relative, rems, Action, Div, FocusHandle, IntoElement, Keystroke};
 
 #[derive(IntoElement, Clone)]
 pub struct KeyBinding {
@@ -49,12 +49,21 @@ impl RenderOnce for KeyBinding {
 
 impl KeyBinding {
     pub fn for_action(action: &dyn Action, cx: &mut WindowContext) -> Option<Self> {
-        // todo! this last is arbitrary, we want to prefer users key bindings over defaults,
-        // and vim over normal (in vim mode), etc.
         let key_binding = cx.bindings_for_action(action).last().cloned()?;
         Some(Self::new(key_binding))
     }
 
+    // like for_action(), but lets you specify the context from which keybindings
+    // are matched.
+    pub fn for_action_in(
+        action: &dyn Action,
+        focus: &FocusHandle,
+        cx: &mut WindowContext,
+    ) -> Option<Self> {
+        let key_binding = cx.bindings_for_action_in(action, focus).last().cloned()?;
+        Some(Self::new(key_binding))
+    }
+
     fn icon_for_key(keystroke: &Keystroke) -> Option<Icon> {
         let mut icon: Option<Icon> = None;