Make remote agent task refs followable

Author: i386

crates/mesh-agents-cli/src/a2a.rs

@@ -270,7 +270,7 @@ impl LocalA2aTools {
                 .with_context(|| {
                     format!("remote task `{task_id}` was not found for agent `{agent_id}`")
                 })?;
-            return serde_json::from_value(remote.result)
+            return decode_remote_task(remote.result)
                 .with_context(|| format!("failed to decode remote task `{task_id}`"));
         }
         let path = agent_task_store_path(&self.data_dir, agent_id);
@@ -313,6 +313,13 @@ impl LocalA2aTools {
     }
 }
 
+fn decode_remote_task(value: Value) -> Result<Task> {
+    if let Some(task) = value.get("task") {
+        return serde_json::from_value(task.clone()).context("failed to decode wrapped task");
+    }
+    serde_json::from_value(value).context("failed to decode task")
+}
+
 #[derive(Clone, Debug, Default)]
 struct AgentGates {
     gates: Arc<Mutex<std::collections::HashMap<String, Arc<AgentGate>>>>,

crates/mesh-agents-cli/src/mesh.rs

@@ -137,24 +137,46 @@ impl RemoteTaskCache {
     }
 
     pub(crate) fn upsert(&mut self, task: RemoteTaskRecord) {
-        if let Some(existing) = self
-            .tasks
-            .iter_mut()
-            .find(|existing| existing.agent_id == task.agent_id && existing.task_id == task.task_id)
-        {
+        if let Some(existing) = self.tasks.iter_mut().find(|existing| {
+            existing.agent_id == task.agent_id
+                && (existing.task_id == task.task_id
+                    || same_nonempty(&existing.correlation_id, &task.correlation_id))
+        }) {
             *existing = task;
         } else {
             self.tasks.push(task);
         }
     }
 
+    pub(crate) fn upsert_pending(
+        &mut self,
+        agent_id: String,
+        peer_id: String,
+        correlation_id: String,
+        result: Value,
+    ) {
+        self.upsert(RemoteTaskRecord {
+            agent_id,
+            task_id: correlation_id.clone(),
+            peer_id,
+            correlation_id,
+            result,
+            updated_at_ms: now_ms(),
+        });
+    }
+
     pub(crate) fn get(&self, agent_id: &str, task_id: &str) -> Option<&RemoteTaskRecord> {
-        self.tasks
-            .iter()
-            .find(|task| task.agent_id == agent_id && task.task_id == task_id)
+        self.tasks.iter().find(|task| {
+            task.agent_id == agent_id
+                && (task.task_id == task_id || same_nonempty(&task.correlation_id, task_id))
+        })
     }
 }
 
+fn same_nonempty(left: &str, right: &str) -> bool {
+    !left.is_empty() && left == right
+}
+
 pub(crate) fn local_advertisements(
     agents_dir: &Path,
     source_peer_id: &str,
@@ -321,4 +343,53 @@ public_mesh = false
         assert_eq!(summaries[0]["agent_id"], "private");
         assert_eq!(summaries[0]["location"], "local");
     }
+
+    #[test]
+    fn remote_tasks_can_be_read_by_correlation_id() {
+        let mut cache = RemoteTaskCache::default();
+        cache.upsert(RemoteTaskRecord {
+            agent_id: "pr-review".to_string(),
+            task_id: "task-1".to_string(),
+            peer_id: "peer-a".to_string(),
+            correlation_id: "remote-1".to_string(),
+            result: json!({ "task": { "id": "task-1" } }),
+            updated_at_ms: 1,
+        });
+
+        let task = cache
+            .get("pr-review", "remote-1")
+            .expect("correlation id should resolve");
+
+        assert_eq!(task.task_id, "task-1");
+    }
+
+    #[test]
+    fn remote_task_response_replaces_pending_correlation_record() {
+        let mut cache = RemoteTaskCache::default();
+        cache.upsert_pending(
+            "pr-review".to_string(),
+            "peer-a".to_string(),
+            "remote-1".to_string(),
+            json!({ "task": { "id": "remote-1" } }),
+        );
+
+        cache.upsert(RemoteTaskRecord {
+            agent_id: "pr-review".to_string(),
+            task_id: "task-1".to_string(),
+            peer_id: "peer-a".to_string(),
+            correlation_id: "remote-1".to_string(),
+            result: json!({ "task": { "id": "task-1" } }),
+            updated_at_ms: 2,
+        });
+
+        assert_eq!(cache.tasks.len(), 1);
+        assert_eq!(
+            cache.get("pr-review", "remote-1").unwrap().task_id,
+            "task-1"
+        );
+        assert_eq!(
+            cache.get("pr-review", "task-1").unwrap().correlation_id,
+            "remote-1"
+        );
+    }
 }

crates/mesh-agents-cli/src/plugin.rs

@@ -422,13 +422,24 @@ async fn plugin_send_message(
     )?;
     message.correlation_id = correlation_id.clone();
     context.send_channel_message(message).await?;
+    let result = pending_remote_task_result(&args.agent_id, &correlation_id, &remote.peer_id);
+    let mut cache = RemoteTaskCache::load(&state.data_dir)?;
+    cache.upsert_pending(
+        args.agent_id.clone(),
+        remote.peer_id.clone(),
+        correlation_id.clone(),
+        result.clone(),
+    );
+    cache.save(&state.data_dir)?;
     Ok(json!({
         "agent_id": args.agent_id,
         "location": "remote",
         "peer_id": remote.peer_id,
+        "task_id": correlation_id.clone(),
         "correlation_id": correlation_id,
         "status": "submitted",
-        "note": "remote task result will be available through agents.get_task after the owner node replies",
+        "result": result,
+        "note": "task_id is a followable pending task reference; agents.get_task will resolve it to the final A2A task after the owner node replies",
     }))
 }
 
@@ -449,6 +460,8 @@ async fn plugin_get_task(state: &PluginState, agent_id: &str, task_id: &str) ->
         "agent_id": agent_id,
         "location": "remote",
         "peer_id": remote.peer_id,
+        "task_id": remote.task_id,
+        "correlation_id": remote.correlation_id,
         "task": remote.result,
     }))
 }
@@ -530,13 +543,32 @@ fn decode_task_value(value: Value) -> Result<Task> {
     serde_json::from_value(value).context("failed to decode task")
 }
 
+fn pending_remote_task_result(agent_id: &str, correlation_id: &str, peer_id: &str) -> Value {
+    json!({
+        "task": {
+            "id": correlation_id,
+            "contextId": correlation_id,
+            "status": {
+                "state": "TASK_STATE_SUBMITTED",
+            },
+            "metadata": {
+                "mesh_llm_pending_remote": true,
+                "agent_id": agent_id,
+                "peer_id": peer_id,
+                "correlation_id": correlation_id,
+            },
+        },
+    })
+}
+
 async fn send_remote_message_stream(
     state: &PluginState,
     remote: &crate::mesh::RemoteAgentAd,
     args: &SendMessageArgs,
     context: &mut PluginContext<'_>,
 ) -> Result<Value> {
     let stream_id = format!("a2a-{}-{}", args.agent_id, now_ms());
+    let correlation_id = format!("remote-{}", now_ms());
     let metadata = MeshStreamRequest::SendMessage {
         agent_id: args.agent_id.clone(),
         message: args.message.clone(),
@@ -552,7 +584,7 @@ async fn send_remote_message_stream(
             mode: mesh_llm_plugin::proto::StreamMode::EventStream as i32,
             bidirectional: false,
             content_type: Some("text/event-stream".to_string()),
-            correlation_id: Some(format!("remote-{}", now_ms())),
+            correlation_id: Some(correlation_id.clone()),
             metadata_json: Some(serde_json::to_string(&metadata)?),
             expected_bytes: None,
             idle_timeout_ms: Some(300_000),
@@ -579,7 +611,7 @@ async fn send_remote_message_stream(
             agent_id: args.agent_id.clone(),
             task_id: task_id.clone(),
             peer_id: remote.peer_id.clone(),
-            correlation_id: String::new(),
+            correlation_id: correlation_id.clone(),
             result: final_result.clone(),
             updated_at_ms: now_ms(),
         });
@@ -592,6 +624,7 @@ async fn send_remote_message_stream(
         "peer_id": remote.peer_id,
         "transport": "mesh_stream",
         "task_id": task_id,
+        "correlation_id": correlation_id,
         "result": final_result,
         "events": events,
     }))