Add a plugin API for reading HTTP response bodies (#560)

crates/yaak/src/plugin_events.rs

@@ -1,3 +1,6 @@
+use crate::response_body::ResponseBodyStore;
+use base64::Engine;
+use base64::prelude::BASE64_STANDARD;
 use yaak_models::models::AnyModel;
 use yaak_models::query_manager::QueryManager;
 use yaak_models::util::UpdateSource;
@@ -5,12 +8,14 @@ use yaak_plugins::events::{
     CloseWindowRequest, CopyTextRequest, DeleteKeyValueRequest, DeleteKeyValueResponse,
     DeleteModelRequest, DeleteModelResponse, ErrorResponse, FindHttpResponsesRequest,
     FindHttpResponsesResponse, GetCookieValueRequest, GetHttpRequestByIdRequest,
-    GetHttpRequestByIdResponse, GetKeyValueRequest, GetKeyValueResponse, InternalEventPayload,
-    ListCookieNamesRequest, ListFoldersRequest, ListFoldersResponse, ListHttpRequestsRequest,
-    ListHttpRequestsResponse, ListOpenWorkspacesRequest, OpenExternalUrlRequest, OpenWindowRequest,
-    PromptFormRequest, PromptTextRequest, ReloadResponse, RenderGrpcRequestRequest,
-    RenderHttpRequestRequest, SendHttpRequestRequest, SetKeyValueRequest, ShowToastRequest,
-    TemplateRenderRequest, UpsertModelRequest, UpsertModelResponse, WindowInfoRequest,
+    GetHttpRequestByIdResponse, GetHttpResponseBodyInfoRequest, GetHttpResponseBodyInfoResponse,
+    GetKeyValueRequest, GetKeyValueResponse, InternalEventPayload, ListCookieNamesRequest,
+    ListFoldersRequest, ListFoldersResponse, ListHttpRequestsRequest, ListHttpRequestsResponse,
+    ListOpenWorkspacesRequest, OpenExternalUrlRequest, OpenWindowRequest, PromptFormRequest,
+    PromptTextRequest, ReadHttpResponseBodyChunkRequest, ReadHttpResponseBodyChunkResponse,
+    ReloadResponse, RenderGrpcRequestRequest, RenderHttpRequestRequest, SendHttpRequestRequest,
+    SetKeyValueRequest, ShowToastRequest, TemplateRenderRequest, UpsertModelRequest,
+    UpsertModelResponse, WindowInfoRequest,
 };
 
 pub struct SharedPluginEventContext<'a> {
@@ -40,6 +45,8 @@ pub enum SharedRequest<'a> {
     ListFolders(&'a ListFoldersRequest),
     ListHttpRequests(&'a ListHttpRequestsRequest),
     FindHttpResponses(&'a FindHttpResponsesRequest),
+    GetHttpResponseBodyInfo(&'a GetHttpResponseBodyInfoRequest),
+    ReadHttpResponseBodyChunk(&'a ReadHttpResponseBodyChunkRequest),
     UpsertModel(&'a UpsertModelRequest),
     DeleteModel(&'a DeleteModelRequest),
 }
@@ -136,6 +143,12 @@ impl<'a> From<&'a InternalEventPayload> for GroupedPluginRequest<'a> {
             InternalEventPayload::FindHttpResponsesRequest(req) => {
                 GroupedPluginRequest::Shared(SharedRequest::FindHttpResponses(req))
             }
+            InternalEventPayload::GetHttpResponseBodyInfoRequest(req) => {
+                GroupedPluginRequest::Shared(SharedRequest::GetHttpResponseBodyInfo(req))
+            }
+            InternalEventPayload::ReadHttpResponseBodyChunkRequest(req) => {
+                GroupedPluginRequest::Shared(SharedRequest::ReadHttpResponseBodyChunk(req))
+            }
             InternalEventPayload::UpsertModelRequest(req) => {
                 GroupedPluginRequest::Shared(SharedRequest::UpsertModel(req))
             }
@@ -182,13 +195,17 @@ impl<'a> From<&'a InternalEventPayload> for GroupedPluginRequest<'a> {
 
 pub fn handle_shared_plugin_event<'a>(
     query_manager: &QueryManager,
+    body_store: &dyn ResponseBodyStore,
     payload: &'a InternalEventPayload,
     context: SharedPluginEventContext<'_>,
 ) -> GroupedPluginEvent<'a> {
     match GroupedPluginRequest::from(payload) {
-        GroupedPluginRequest::Shared(req) => {
-            GroupedPluginEvent::Handled(Some(build_shared_reply(query_manager, req, context)))
-        }
+        GroupedPluginRequest::Shared(req) => GroupedPluginEvent::Handled(Some(build_shared_reply(
+            query_manager,
+            body_store,
+            req,
+            context,
+        ))),
         GroupedPluginRequest::Host(req) => GroupedPluginEvent::ToHandle(req),
         GroupedPluginRequest::Ignore => GroupedPluginEvent::Handled(None),
     }
@@ -196,6 +213,7 @@ pub fn handle_shared_plugin_event<'a>(
 
 fn build_shared_reply(
     query_manager: &QueryManager,
+    body_store: &dyn ResponseBodyStore,
     request: SharedRequest<'_>,
     context: SharedPluginEventContext<'_>,
 ) -> InternalEventPayload {
@@ -283,6 +301,31 @@ fn build_shared_reply(
                 http_responses,
             })
         }
+        SharedRequest::GetHttpResponseBodyInfo(req) => match body_store.info(&req.response_id) {
+            Ok(info) => InternalEventPayload::GetHttpResponseBodyInfoResponse(
+                GetHttpResponseBodyInfoResponse {
+                    content_length: info.content_length,
+                    content_type: info.content_type,
+                    complete: info.complete,
+                },
+            ),
+            Err(err) => InternalEventPayload::ErrorResponse(ErrorResponse {
+                error: format!("Failed to read body of response {}: {err}", req.response_id),
+            }),
+        },
+        SharedRequest::ReadHttpResponseBodyChunk(req) => {
+            match body_store.read_chunk(&req.response_id, req.offset, req.length) {
+                Ok(bytes) => InternalEventPayload::ReadHttpResponseBodyChunkResponse(
+                    ReadHttpResponseBodyChunkResponse {
+                        length: bytes.len() as u64,
+                        data: BASE64_STANDARD.encode(bytes),
+                    },
+                ),
+                Err(err) => InternalEventPayload::ErrorResponse(ErrorResponse {
+                    error: format!("Failed to read body of response {}: {err}", req.response_id),
+                }),
+            }
+        }
         SharedRequest::UpsertModel(req) => {
             use AnyModel::*;
 
@@ -437,10 +480,26 @@ fn build_shared_reply(
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::response_body::{FileResponseBodyStore, ResponseBodyInfo};
+    use std::cell::RefCell;
     use tempfile::TempDir;
     use yaak_models::models::{AnyModel, Folder, HttpRequest, Workspace};
     use yaak_models::util::UpdateSource;
 
+    /// The real dispatch, with the store the desktop and CLI hand it.
+    fn dispatch<'a>(
+        query_manager: &QueryManager,
+        payload: &'a InternalEventPayload,
+        context: SharedPluginEventContext<'_>,
+    ) -> GroupedPluginEvent<'a> {
+        handle_shared_plugin_event(
+            query_manager,
+            &FileResponseBodyStore::new(query_manager),
+            payload,
+            context,
+        )
+    }
+
     fn seed_query_manager() -> (QueryManager, TempDir) {
         let temp_dir = TempDir::new().expect("Failed to create temp dir");
         let db_path = temp_dir.path().join("db.sqlite");
@@ -498,7 +557,7 @@ mod tests {
         let payload = InternalEventPayload::ListHttpRequestsRequest(
             yaak_plugins::events::ListHttpRequestsRequest { folder_id: None },
         );
-        let result = handle_shared_plugin_event(
+        let result = dispatch(
             &query_manager,
             &payload,
             SharedPluginEventContext { plugin_name: "@yaak/test", workspace_id: None },
@@ -517,7 +576,7 @@ mod tests {
         let by_workspace_payload = InternalEventPayload::ListHttpRequestsRequest(
             yaak_plugins::events::ListHttpRequestsRequest { folder_id: None },
         );
-        let by_workspace = handle_shared_plugin_event(
+        let by_workspace = dispatch(
             &query_manager,
             &by_workspace_payload,
             SharedPluginEventContext { plugin_name: "@yaak/test", workspace_id: Some("wk_test") },
@@ -536,7 +595,7 @@ mod tests {
                 folder_id: Some("fl_test".to_string()),
             },
         );
-        let by_folder = handle_shared_plugin_event(
+        let by_folder = dispatch(
             &query_manager,
             &by_folder_payload,
             SharedPluginEventContext { plugin_name: "@yaak/test", workspace_id: None },
@@ -559,7 +618,7 @@ mod tests {
             limit: Some(1),
         });
 
-        let result = handle_shared_plugin_event(
+        let result = dispatch(
             &query_manager,
             &payload,
             SharedPluginEventContext { plugin_name: "@yaak/test", workspace_id: Some("wk_test") },
@@ -575,6 +634,105 @@ mod tests {
         }
     }
 
+    /// A store that answers from memory, standing in for whatever holds the
+    /// bytes — the point being that the dispatch below never learns which.
+    struct FakeBodyStore {
+        body: Vec<u8>,
+        reads: RefCell<Vec<(u64, u64)>>,
+    }
+
+    impl ResponseBodyStore for FakeBodyStore {
+        fn info(&self, _response_id: &str) -> crate::error::Result<ResponseBodyInfo> {
+            Ok(ResponseBodyInfo {
+                content_length: self.body.len() as u64,
+                content_type: Some("text/plain; charset=utf-8".to_string()),
+                complete: true,
+            })
+        }
+
+        fn read_chunk(
+            &self,
+            _response_id: &str,
+            offset: u64,
+            length: u64,
+        ) -> crate::error::Result<Vec<u8>> {
+            self.reads.borrow_mut().push((offset, length));
+            let start = (offset as usize).min(self.body.len());
+            let end = (start + length as usize).min(self.body.len());
+            Ok(self.body[start..end].to_vec())
+        }
+    }
+
+    #[test]
+    fn response_body_is_read_by_id_through_the_store() {
+        let (query_manager, _temp_dir) = seed_query_manager();
+        let store = FakeBodyStore { body: b"hello".to_vec(), reads: RefCell::new(Vec::new()) };
+
+        let info_payload = InternalEventPayload::GetHttpResponseBodyInfoRequest(
+            GetHttpResponseBodyInfoRequest { response_id: "rs_test".to_string() },
+        );
+        let info = handle_shared_plugin_event(
+            &query_manager,
+            &store,
+            &info_payload,
+            SharedPluginEventContext { plugin_name: "@yaak/test", workspace_id: None },
+        );
+        match info {
+            GroupedPluginEvent::Handled(Some(
+                InternalEventPayload::GetHttpResponseBodyInfoResponse(resp),
+            )) => {
+                assert_eq!(resp.content_length, 5);
+                assert_eq!(resp.content_type.as_deref(), Some("text/plain; charset=utf-8"));
+            }
+            other => panic!("unexpected body info result: {other:?}"),
+        }
+
+        let chunk_payload = InternalEventPayload::ReadHttpResponseBodyChunkRequest(
+            ReadHttpResponseBodyChunkRequest {
+                response_id: "rs_test".to_string(),
+                offset: 1,
+                length: 3,
+            },
+        );
+        let chunk = handle_shared_plugin_event(
+            &query_manager,
+            &store,
+            &chunk_payload,
+            SharedPluginEventContext { plugin_name: "@yaak/test", workspace_id: None },
+        );
+        match chunk {
+            GroupedPluginEvent::Handled(Some(
+                InternalEventPayload::ReadHttpResponseBodyChunkResponse(resp),
+            )) => {
+                assert_eq!(resp.length, 3);
+                assert_eq!(BASE64_STANDARD.decode(resp.data).unwrap(), b"ell");
+            }
+            other => panic!("unexpected body chunk result: {other:?}"),
+        }
+
+        assert_eq!(*store.reads.borrow(), vec![(1, 3)]);
+    }
+
+    #[test]
+    fn an_unreadable_response_body_becomes_an_error_reply() {
+        let (query_manager, _temp_dir) = seed_query_manager();
+        let payload = InternalEventPayload::GetHttpResponseBodyInfoRequest(
+            GetHttpResponseBodyInfoRequest { response_id: "rs_never_persisted".to_string() },
+        );
+        let result = dispatch(
+            &query_manager,
+            &payload,
+            SharedPluginEventContext { plugin_name: "@yaak/test", workspace_id: None },
+        );
+
+        match result {
+            GroupedPluginEvent::Handled(Some(InternalEventPayload::ErrorResponse(resp))) => {
+                assert!(resp.error.contains("rs_never_persisted"), "unhelpful error: {}", resp.error)
+            }
+            other => panic!("unexpected missing-response result: {other:?}"),
+        }
+    }
+
     #[test]
     fn upsert_and_delete_model_are_shared_handled() {
         let (query_manager, _temp_dir) = seed_query_manager();
@@ -590,7 +748,7 @@ mod tests {
             }),
         });
 
-        let upsert_result = handle_shared_plugin_event(
+        let upsert_result = dispatch(
             &query_manager,
             &upsert_payload,
             SharedPluginEventContext { plugin_name: "@yaak/test", workspace_id: Some("wk_test") },
@@ -609,7 +767,7 @@ mod tests {
             model: "http_request".to_string(),
             id: "rq_test".to_string(),
         });
-        let delete_result = handle_shared_plugin_event(
+        let delete_result = dispatch(
             &query_manager,
             &delete_payload,
             SharedPluginEventContext { plugin_name: "@yaak/test", workspace_id: Some("wk_test") },
@@ -631,7 +789,7 @@ mod tests {
         let payload = InternalEventPayload::WindowInfoRequest(WindowInfoRequest {
             label: "main".to_string(),
         });
-        let result = handle_shared_plugin_event(
+        let result = dispatch(
             &query_manager,
             &payload,
             SharedPluginEventContext { plugin_name: "@yaak/test", workspace_id: None },