diff --git a/api/libs/handlers/src/chats/post_chat_handler.rs b/api/libs/handlers/src/chats/post_chat_handler.rs
index f21c5a1e9..1bb5d1c1f 100644
--- a/api/libs/handlers/src/chats/post_chat_handler.rs
+++ b/api/libs/handlers/src/chats/post_chat_handler.rs
@@ -235,24 +235,40 @@ pub async fn post_chat_handler(
// Only store completed messages in raw_llm_messages
match &msg {
AgentMessage::Assistant {
- progress, content, ..
+ progress, content, id, ..
} => {
- if let Some(content) = content {
- raw_response_message.push_str(&content);
+ // Store chunks in the tracker to ensure deduplication
+ if let Some(content_str) = content {
+ // Use message ID as chunk ID, or generate a consistent one if missing
+ let chunk_id = id.clone().unwrap_or_else(|| "assistant_message".to_string());
+ // Add to chunk tracker to handle deduplication
+ chunk_tracker.add_chunk(chunk_id.clone(), content_str.clone());
}
if matches!(progress, MessageProgress::Complete) {
- if raw_response_message.is_empty() {
- raw_llm_messages.push(msg.clone());
- } else {
+ if let Some(content_str) = content {
+ // Use message ID as chunk ID, or generate a consistent one if missing
+ let chunk_id = id.clone().unwrap_or_else(|| "assistant_message".to_string());
+
+ // Get the complete deduplicated text from the chunk tracker
+ let complete_text = chunk_tracker.get_complete_text(chunk_id.clone())
+ .unwrap_or_else(|| content_str.clone());
+
+ // Create a new message with the deduplicated content
raw_llm_messages.push(AgentMessage::Assistant {
- id: None,
- content: Some(raw_response_message.clone()),
+ id: id.clone(),
+ content: Some(complete_text),
name: None,
tool_calls: None,
progress: MessageProgress::Complete,
initial: false,
});
+
+ // Clear the chunk from the tracker
+ chunk_tracker.clear_chunk(chunk_id);
+ } else {
+ // If there's no content, just use the original message
+ raw_llm_messages.push(msg.clone());
}
}
}
diff --git a/api/libs/sharing/src/remove_asset_permissions.rs b/api/libs/sharing/src/remove_asset_permissions.rs
index d31321a19..90e4c1e88 100644
--- a/api/libs/sharing/src/remove_asset_permissions.rs
+++ b/api/libs/sharing/src/remove_asset_permissions.rs
@@ -1,7 +1,7 @@
use anyhow::{Context, Result};
use chrono::Utc;
use database::{
- enums::{AssetType, IdentityType},
+ enums::{AssetPermissionRole, AssetType, IdentityType},
pool::get_pg_pool,
schema::asset_permissions,
};
@@ -9,7 +9,7 @@ use diesel::ExpressionMethods;
use diesel_async::RunQueryDsl;
use uuid::Uuid;
-use crate::{errors::SharingError, user_lookup::find_user_by_email};
+use crate::{check_access, errors::SharingError, user_lookup::find_user_by_email};
/// Removes a sharing record for a specific user + asset combination
pub async fn remove_share(
@@ -55,6 +55,8 @@ pub async fn remove_share(
}
/// Removes a sharing record identified by user email
+///
+/// Requires the caller to have Owner or FullAccess permission on the asset.
pub async fn remove_share_by_email(
email: &str,
asset_id: Uuid,
@@ -66,6 +68,24 @@ pub async fn remove_share_by_email(
return Err(SharingError::InvalidEmail(email.to_string()).into());
}
+ // Validate that the caller has permission to remove shares (Owner or FullAccess)
+ let caller_role = check_access(asset_id, asset_type, updated_by, IdentityType::User).await?;
+
+ match caller_role {
+ Some(AssetPermissionRole::Owner) | Some(AssetPermissionRole::FullAccess) => {
+ // Caller has permission to remove shares
+ },
+ Some(_) => {
+ return Err(SharingError::InsufficientPermissions.into());
+ },
+ None => {
+ return Err(SharingError::PermissionDenied(format!(
+ "User does not have permission to access asset {} of type {:?}",
+ asset_id, asset_type
+ )).into());
+ }
+ }
+
// Find the user by email
let user = find_user_by_email(email)
.await?
@@ -78,7 +98,7 @@ pub async fn remove_share_by_email(
#[cfg(test)]
mod tests {
use super::*;
- use database::enums::{AssetType, IdentityType};
+ use database::enums::{AssetPermissionRole, AssetType, IdentityType};
use uuid::Uuid;
#[test]
@@ -95,4 +115,17 @@ mod tests {
let err = result.unwrap_err().to_string();
assert!(err.contains("Invalid email"));
}
+
+ // Note: The following tests are commented out as they would require
+ // more complex integration test setup with database access.
+ // In a real implementation, we would use integration tests to cover these cases.
+
+ // Test case: User has CanEdit permission but not Owner or FullAccess
+ // Expected result: InsufficientPermissions error
+
+ // Test case: User not found by email
+ // Expected result: UserNotFound error
+
+ // Test case: Successful removal when user has Owner permission
+ // Expected result: Success
}
\ No newline at end of file
diff --git a/api/libs/sharing/src/user_lookup.rs b/api/libs/sharing/src/user_lookup.rs
index 3e1b9a7e0..429b3b3de 100644
--- a/api/libs/sharing/src/user_lookup.rs
+++ b/api/libs/sharing/src/user_lookup.rs
@@ -12,7 +12,7 @@ use crate::errors::SharingError;
///
/// # Returns
/// * `Ok(Some(User))` - If the user was found
-/// * `Ok(None)` - If no user with that email exists
+/// * `Ok(None)` - If no user with that email exists or if the user has been deleted
/// * `Err(_)` - If there was a database error or other issue
///
pub async fn find_user_by_email(email: &str) -> Result