mirror of
https://github.com/mountain-loop/yaak.git
synced 2026-09-16 14:51:37 +02:00
Fix MCP send_http_request ignoring environmentId (#654)
This commit is contained in:
Ricardo Lopes Temoteo
@@ -182,6 +182,31 @@ async fn build_plugin_reply(
|
||||
http_request.workspace_id = workspace_id;
|
||||
}
|
||||
|
||||
let environment_id = if let Some(environment_id) =
|
||||
send_http_request_request.environment_id.as_deref()
|
||||
{
|
||||
if shared_workspace_id.is_some_and(|id| id != http_request.workspace_id) {
|
||||
return Some(InternalEventPayload::ErrorResponse(ErrorResponse {
|
||||
error: "HTTP request does not belong to the selected workspace"
|
||||
.to_string(),
|
||||
}));
|
||||
}
|
||||
match host_context
|
||||
.query_manager
|
||||
.connect()
|
||||
.get_environment_for_workspace(&http_request.workspace_id, environment_id)
|
||||
{
|
||||
Ok(environment) => Some(environment.id),
|
||||
Err(err) => {
|
||||
return Some(InternalEventPayload::ErrorResponse(ErrorResponse {
|
||||
error: err.to_string(),
|
||||
}));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
execution_context.environment_id.clone()
|
||||
};
|
||||
|
||||
let cookie_jar_id =
|
||||
if let Some(cookie_jar_id) = execution_context.cookie_jar_id.clone() {
|
||||
Some(cookie_jar_id)
|
||||
@@ -211,7 +236,7 @@ async fn build_plugin_reply(
|
||||
query_manager: &host_context.query_manager,
|
||||
blob_manager: &host_context.blob_manager,
|
||||
request: http_request,
|
||||
environment_id: execution_context.environment_id.as_deref(),
|
||||
environment_id: environment_id.as_deref(),
|
||||
update_source: UpdateSource::Plugin,
|
||||
cookie_jar_id,
|
||||
response_dir: &host_context.response_dir,
|
||||
@@ -1003,3 +1028,180 @@ fn prompt_label_for_base(base: &yaak_plugins::events::FormInputBase) -> String {
|
||||
}
|
||||
base.name.clone()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod environment_tests {
|
||||
use super::*;
|
||||
use serde_json::json;
|
||||
use tempfile::TempDir;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::time::{Duration, timeout};
|
||||
use yaak_models::models::{EnvironmentVariable, HttpRequest, HttpRequestHeader};
|
||||
|
||||
#[tokio::test]
|
||||
async fn plugin_send_uses_override_and_fallback_and_rejects_invalid_ids_before_sending() {
|
||||
let dir = TempDir::new().unwrap();
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let (received_tx, mut received_rx) = tokio::sync::mpsc::unbounded_channel();
|
||||
let server = tokio::spawn(async move {
|
||||
loop {
|
||||
let (mut socket, _) = listener.accept().await.unwrap();
|
||||
let mut request = Vec::new();
|
||||
loop {
|
||||
let mut chunk = [0u8; 1024];
|
||||
let n = socket.read(&mut chunk).await.unwrap();
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
request.extend_from_slice(&chunk[..n]);
|
||||
if request.windows(4).any(|w| w == b"\r\n\r\n") {
|
||||
break;
|
||||
}
|
||||
}
|
||||
let request = String::from_utf8(request).unwrap();
|
||||
eprintln!("Local HTTP test received on {address}:\n{request}");
|
||||
received_tx.send(request).unwrap();
|
||||
socket
|
||||
.write_all(
|
||||
b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok",
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
});
|
||||
let (query_manager, blob_manager, _rx) = yaak_models::init_standalone(
|
||||
&dir.path().join("db.sqlite"),
|
||||
&dir.path().join("blobs.sqlite"),
|
||||
)
|
||||
.unwrap();
|
||||
let (base, a, b, request) = query_manager
|
||||
.with_tx(|db| {
|
||||
let source = &UpdateSource::Background;
|
||||
let workspace = db.list_workspaces()?.remove(0);
|
||||
let vars = |value: &str| {
|
||||
vec![EnvironmentVariable {
|
||||
enabled: true,
|
||||
name: "marker".into(),
|
||||
value: value.into(),
|
||||
..Default::default()
|
||||
}]
|
||||
};
|
||||
let base = db.ensure_base_environment(&workspace.id)?;
|
||||
let base = db.upsert_environment(
|
||||
&Environment { variables: vars("global"), ..base },
|
||||
source,
|
||||
)?;
|
||||
let sub = |name: &str| {
|
||||
db.upsert_environment(
|
||||
&Environment {
|
||||
name: name.into(),
|
||||
workspace_id: workspace.id.clone(),
|
||||
parent_model: "environment".into(),
|
||||
parent_id: Some(base.id.clone()),
|
||||
variables: vars(name),
|
||||
..Default::default()
|
||||
},
|
||||
source,
|
||||
)
|
||||
};
|
||||
let a = sub("a")?;
|
||||
let b = sub("b")?;
|
||||
let request = db.upsert_http_request(
|
||||
&HttpRequest {
|
||||
workspace_id: workspace.id,
|
||||
url: format!("http://{address}/echo"),
|
||||
method: "GET".into(),
|
||||
headers: vec![HttpRequestHeader {
|
||||
enabled: true,
|
||||
name: "X-Environment".into(),
|
||||
value: "${[ marker ]}".into(),
|
||||
..Default::default()
|
||||
}],
|
||||
..Default::default()
|
||||
},
|
||||
source,
|
||||
)?;
|
||||
Ok::<_, yaak_models::error::Error>((base, a, b, request))
|
||||
})
|
||||
.unwrap();
|
||||
let plugin_dir = dir.path().join("plugins");
|
||||
std::fs::create_dir_all(&plugin_dir).unwrap();
|
||||
let plugin_manager = Arc::new(
|
||||
PluginManager::new(
|
||||
plugin_dir.clone(),
|
||||
plugin_dir,
|
||||
PathBuf::from("node"),
|
||||
PathBuf::from(env!("CARGO_MANIFEST_DIR"))
|
||||
.join("../../crates-tauri/yaak-app-client/vendored/plugin-runtime/index.cjs"),
|
||||
&query_manager,
|
||||
&PluginContext::new_empty(),
|
||||
false,
|
||||
)
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
let host = CliHostContext {
|
||||
encryption_manager: Arc::new(EncryptionManager::new(
|
||||
query_manager.clone(),
|
||||
"yaak-test",
|
||||
)),
|
||||
query_manager,
|
||||
blob_manager,
|
||||
plugin_manager: plugin_manager.clone(),
|
||||
connection_manager: Arc::new(HttpConnectionManager::new()),
|
||||
response_dir: dir.path().join("responses"),
|
||||
execution_context: CliExecutionContext {
|
||||
workspace_id: Some(base.workspace_id.clone()),
|
||||
environment_id: Some(a.id.clone()),
|
||||
..Default::default()
|
||||
},
|
||||
};
|
||||
std::fs::create_dir_all(&host.response_dir).unwrap();
|
||||
let event = |environment_id: Option<&str>| -> InternalEvent {
|
||||
let mut payload =
|
||||
json!({ "type": "send_http_request_request", "httpRequest": request });
|
||||
if let Some(id) = environment_id {
|
||||
payload["environmentId"] = json!(id);
|
||||
}
|
||||
serde_json::from_value(json!({
|
||||
"id": "test", "pluginRefId": "test", "pluginName": "test", "replyId": null,
|
||||
"context": PluginContext::new_empty(), "payload": payload
|
||||
}))
|
||||
.unwrap()
|
||||
};
|
||||
for (id, marker) in [
|
||||
(Some(b.id.as_str()), "b"),
|
||||
(None, "a"),
|
||||
(Some(base.id.as_str()), "global"),
|
||||
] {
|
||||
let reply =
|
||||
timeout(Duration::from_secs(10), build_plugin_reply(&host, &event(id), "test"))
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(
|
||||
matches!(reply, Some(InternalEventPayload::SendHttpRequestResponse(ref r)) if r.http_response.status == 200),
|
||||
"{reply:?}"
|
||||
);
|
||||
let received = received_rx.recv().await.unwrap().to_lowercase();
|
||||
assert!(received.contains(&format!("x-environment: {marker}\r\n")), "{received}");
|
||||
assert_eq!(host.execution_context.environment_id.as_deref(), Some(a.id.as_str()));
|
||||
}
|
||||
for id in ["", "ev_missing"] {
|
||||
let reply = build_plugin_reply(&host, &event(Some(id)), "test").await;
|
||||
assert!(matches!(reply, Some(InternalEventPayload::ErrorResponse(_))), "{reply:?}");
|
||||
}
|
||||
assert!(received_rx.try_recv().is_err());
|
||||
assert_eq!(
|
||||
host.query_manager
|
||||
.connect()
|
||||
.list_http_responses_for_request(&request.id, None)
|
||||
.unwrap()
|
||||
.len(),
|
||||
3
|
||||
);
|
||||
plugin_manager.terminate().await;
|
||||
server.abort();
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user