Workspace rules for Claude sessions will use the original session ID if available

Author: krlvi

crates/but-claude/src/db.rs

@@ -52,6 +52,20 @@ pub fn get_session_by_id(
     }
 }
 
+pub fn get_session_by_current_id(
+    ctx: &mut CommandContext,
+    current_id: Uuid,
+) -> anyhow::Result<Option<ClaudeSession>> {
+    let session = ctx
+        .db()?
+        .claude_sessions()
+        .get_by_current_id(&current_id.to_string())?;
+    match session {
+        Some(s) => Ok(Some(s.try_into()?)),
+        None => Ok(None),
+    }
+}
+
 /// Deletes a Claude session and all associated messages from the database. This is what we want to use when we want to delete a session completely.
 pub fn delete_session_and_messages_by_id(
     ctx: &mut CommandContext,

crates/but-claude/src/hooks/mod.rs

@@ -118,10 +118,11 @@ pub async fn handle_stop() -> anyhow::Result<ClaudeHookOutput> {
     let prompt = transcript.prompt().unwrap_or_default();
 
     let ctx = &mut CommandContext::open(&project, AppSettings::load_from_default_path_creating()?)?;
+    let session_id = original_session_id(ctx, input.session_id.clone())?;
 
     let defer = ClearLocksGuard {
         ctx,
-        session_id: input.session_id.clone(),
+        session_id: session_id.clone(),
         file_path: None,
     };
 
@@ -131,14 +132,14 @@ pub async fn handle_stop() -> anyhow::Result<ClaudeHookOutput> {
 
     // If the session stopped, but there's no session persisted in the database, we create a new one.
     // If the session is already persisted, we just retrieve it.
-    let stack_id = get_or_create_session(defer.ctx, &input.session_id, stacks, vb_state)?;
+    let stack_id = get_or_create_session(defer.ctx, &session_id, stacks, vb_state)?;
 
     let (id, outcome) = but_action::handle_changes(
         defer.ctx,
         &summary,
         Some(prompt.clone()),
         ActionHandler::HandleChangesSimple,
-        Source::ClaudeCode(input.session_id),
+        Source::ClaudeCode(session_id),
         Some(stack_id),
     )?;
 
@@ -310,8 +311,9 @@ pub fn handle_pre_tool_call() -> anyhow::Result<ClaudeHookOutput> {
     input.tool_input.file_path = relative_file_path;
 
     let ctx = &mut CommandContext::open(&project, AppSettings::load_from_default_path_creating()?)?;
+    let session_id = original_session_id(ctx, input.session_id.clone())?;
 
-    file_lock::obtain(ctx, input.session_id, input.tool_input.file_path.clone())?;
+    file_lock::obtain(ctx, session_id, input.tool_input.file_path.clone())?;
 
     Ok(ClaudeHookOutput {
         do_continue: true,
@@ -355,18 +357,19 @@ pub fn handle_post_tool_call() -> anyhow::Result<ClaudeHookOutput> {
     input.tool_response.file_path = relative_file_path;
 
     let ctx = &mut CommandContext::open(&project, AppSettings::load_from_default_path_creating()?)?;
+    let session_id = original_session_id(ctx, input.session_id.clone())?;
 
     let defer = ClearLocksGuard {
         ctx,
-        session_id: input.session_id.clone(),
+        session_id: session_id.clone(),
         file_path: Some(input.tool_response.file_path.clone()),
     };
 
     let stacks = list_stacks(defer.ctx)?;
 
     let vb_state = &VirtualBranchesHandle::new(defer.ctx.project().gb_dir());
 
-    let stack_id = get_or_create_session(defer.ctx, &input.session_id, stacks, vb_state)?;
+    let stack_id = get_or_create_session(defer.ctx, &session_id, stacks, vb_state)?;
 
     let changes = but_core::diff::ui::worktree_changes_by_worktree_dir(project.path)?.changes;
     let (assignments, _assignments_error) = but_hunk_assignment::assignments_with_fallback(
@@ -404,6 +407,16 @@ pub fn handle_post_tool_call() -> anyhow::Result<ClaudeHookOutput> {
     })
 }
 
+fn original_session_id(ctx: &mut CommandContext, current_id: String) -> Result<String> {
+    let original_session_id =
+        crate::db::get_session_by_current_id(ctx, Uuid::parse_str(&current_id)?)?;
+    if let Some(session) = original_session_id {
+        Ok(session.id.to_string())
+    } else {
+        Ok(current_id)
+    }
+}
+
 fn get_or_create_session(
     ctx: &mut CommandContext,
     session_id: &str,

crates/but-db/src/claude.rs

@@ -112,6 +112,17 @@ impl ClaudeSessionsHandle<'_> {
         Ok(session)
     }
 
+    pub fn get_by_current_id(
+        &mut self,
+        current_id: &str,
+    ) -> Result<Option<ClaudeSession>, diesel::result::Error> {
+        let session = claude_sessions
+            .filter(crate::schema::claude_sessions::current_id.eq(current_id))
+            .first::<ClaudeSession>(&mut self.db.conn)
+            .optional()?;
+        Ok(session)
+    }
+
     pub fn list(&mut self) -> Result<Vec<ClaudeSession>, diesel::result::Error> {
         let sessions = claude_sessions.load::<ClaudeSession>(&mut self.db.conn)?;
         Ok(sessions)