fix: echo REST completion nonce and argument context

crates/runtime-api/src/lib.rs

@@ -569,12 +569,18 @@ pub struct CompletionRequest {
     pub prefix: String,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub context: Option<protocol::CompletionContext>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub request_id: Option<String>,
 }
 
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
 pub struct CompletionResponse {
     pub kind: CompletionKind,
     pub prefix: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub context: Option<protocol::CompletionContext>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub request_id: Option<String>,
     pub entries: Vec<CompletionEntry>,
 }
 

crates/server-api/src/lib.rs

@@ -8273,6 +8273,8 @@ pub struct RuntimeWorkerCompletionsRequest {
     pub prefix: String,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub context: Option<protocol::CompletionContext>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub request_id: Option<String>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, JsonSchema)]
@@ -8282,6 +8284,10 @@ pub struct RuntimeWorkerCompletionsResult {
     pub worker_id: String,
     pub kind: protocol::CompletionKind,
     pub prefix: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub context: Option<protocol::CompletionContext>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub request_id: Option<String>,
     pub entries: Vec<protocol::CompletionEntry>,
     #[serde(default)]
     pub diagnostics: Vec<Diagnostic>,

crates/worker-runtime/src/http_server.rs

@@ -521,21 +521,8 @@ pub struct RuntimeHttpUploadedFileDeleteResponse {
     pub deleted: bool,
 }
 
-#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
-pub struct RuntimeHttpWorkerCompletionsRequest {
-    pub kind: protocol::CompletionKind,
-    #[serde(default)]
-    pub prefix: String,
-    #[serde(default, skip_serializing_if = "Option::is_none")]
-    pub context: Option<protocol::CompletionContext>,
-}
-
-#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
-pub struct RuntimeHttpWorkerCompletionsResponse {
-    pub kind: protocol::CompletionKind,
-    pub prefix: String,
-    pub entries: Vec<protocol::CompletionEntry>,
-}
+pub type RuntimeHttpWorkerCompletionsRequest = runtime_api::CompletionRequest;
+pub type RuntimeHttpWorkerCompletionsResponse = runtime_api::CompletionResponse;
 
 /// Worker lifecycle request body used by stop/cancel endpoints.
 #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
@@ -1880,6 +1867,8 @@ async fn worker_completions(
     Ok(Json(RuntimeHttpWorkerCompletionsResponse {
         kind: request.kind,
         prefix: request.prefix,
+        context: request.context,
+        request_id: request.request_id,
         entries,
     }))
 }
@@ -3517,6 +3506,56 @@ mod tests {
         serde_json::from_slice(&body).unwrap()
     }
 
+    #[tokio::test]
+    async fn rest_feature_completions_echo_each_nonce_and_qualified_context() {
+        let runtime =
+            Runtime::with_execution_backend(RuntimeOptions::default(), Arc::new(AcceptingBackend))
+                .unwrap();
+        runtime
+            .store_config_bundle(test_bundle(ProfileSelector::Builtin(
+                "builtin:coder".into(),
+            )))
+            .unwrap();
+        let token = "completion-token";
+        let app = runtime_http_router(runtime, token.into());
+        let created = authed_json_request(
+            app.clone(),
+            Method::POST,
+            "/v1/workers",
+            token,
+            &task_request("completion"),
+        )
+        .await;
+        assert_eq!(created.status(), StatusCode::OK);
+        let created: RuntimeHttpWorkerResponse = read_json(created).await;
+        let context = protocol::CompletionContext {
+            invocation: protocol::FeatureInvocationIdentity("builtin:test/prepare".into()),
+            argument: Some("path".into()),
+        };
+        for nonce in ["first", "second", "first"] {
+            let request = runtime_api::CompletionRequest {
+                kind: protocol::CompletionKind::FeatureArgument,
+                prefix: "資料/".into(),
+                context: Some(context.clone()),
+                request_id: Some(nonce.into()),
+            };
+            let response = authed_json_request(
+                app.clone(),
+                Method::POST,
+                &format!("/v1/workers/{}/completions", created.worker.worker_id),
+                token,
+                &request,
+            )
+            .await;
+            assert_eq!(response.status(), StatusCode::OK);
+            let response: runtime_api::CompletionResponse = read_json(response).await;
+            assert_eq!(response.kind, request.kind);
+            assert_eq!(response.prefix, request.prefix);
+            assert_eq!(response.context, request.context);
+            assert_eq!(response.request_id, request.request_id);
+        }
+    }
+
     #[tokio::test]
     async fn rest_command_api_delegates_to_runtime() {
         let runtime =

crates/worker-runtime/src/http_server/runtime_management_api.rs

@@ -582,6 +582,8 @@ impl runtime_api::RuntimeApi for RuntimeManagementApi {
         Ok(runtime_api::CompletionResponse {
             kind,
             prefix,
+            context: value.context,
+            request_id: value.request_id,
             entries,
         })
     }

crates/workspace-server/src/hosts.rs

@@ -737,6 +737,8 @@ pub struct WorkerCompletionsRequest {
     pub prefix: String,
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub context: Option<protocol::CompletionContext>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub request_id: Option<String>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
@@ -745,6 +747,10 @@ pub struct WorkerCompletionsResult {
     pub worker: RuntimeWorkerRef,
     pub kind: protocol::CompletionKind,
     pub prefix: String,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub context: Option<protocol::CompletionContext>,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub request_id: Option<String>,
     pub entries: Vec<protocol::CompletionEntry>,
     pub diagnostics: Vec<RuntimeDiagnostic>,
 }
@@ -1329,6 +1335,8 @@ pub trait WorkspaceWorkerRuntime: Send + Sync {
             worker: RuntimeWorkerRef::new(self.runtime_id().to_string(), worker_id.to_string()),
             kind: request.kind,
             prefix: request.prefix,
+            context: request.context,
+            request_id: request.request_id,
             entries: Vec::new(),
             diagnostics: vec![diagnostic(
                 "worker_completions_unsupported",
@@ -3312,6 +3320,8 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
                 worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
                 kind: request.kind,
                 prefix: request.prefix,
+                context: request.context,
+                request_id: request.request_id,
                 entries: Vec::new(),
                 diagnostics: vec![diagnostic(
                     "embedded_worker_execution_unavailable",
@@ -3327,6 +3337,8 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
                 worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
                 kind: request.kind,
                 prefix: request.prefix,
+                context: request.context,
+                request_id: request.request_id,
                 entries: Vec::new(),
                 diagnostics: vec![diagnostic(
                     "embedded_worker_id_invalid",
@@ -3345,6 +3357,8 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
                 worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
                 kind: request.kind,
                 prefix: request.prefix,
+                context: request.context,
+                request_id: request.request_id,
                 entries,
                 diagnostics: Vec::new(),
             },
@@ -3352,6 +3366,8 @@ impl WorkspaceWorkerRuntime for EmbeddedWorkerRuntime {
                 worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
                 kind: request.kind,
                 prefix: request.prefix,
+                context: request.context,
+                request_id: request.request_id,
                 entries: Vec::new(),
                 diagnostics: vec![embedded_runtime_diagnostic(&error)],
             },
@@ -5420,9 +5436,12 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
             kind: request.kind,
             prefix: request.prefix,
             context: request.context,
+            request_id: request.request_id,
         };
         let failure_kind = request.kind;
         let failure_prefix = request.prefix.clone();
+        let failure_context = request.context.clone();
+        let failure_request_id = request.request_id.clone();
         let worker_id_owned = worker_id.to_string();
         match self.run_runtime_api(
             self.request_timeout,
@@ -5437,6 +5456,8 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
                 worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
                 kind: response.kind,
                 prefix: response.prefix,
+                context: response.context,
+                request_id: response.request_id,
                 entries: response.entries,
                 diagnostics: Vec::new(),
             },
@@ -5444,6 +5465,8 @@ impl WorkspaceWorkerRuntime for RemoteWorkerRuntime {
                 worker: RuntimeWorkerRef::new(self.runtime_id.clone(), worker_id.to_string()),
                 kind: failure_kind,
                 prefix: failure_prefix,
+                context: failure_context,
+                request_id: failure_request_id,
                 entries: Vec::new(),
                 diagnostics: vec![diagnostic],
             },
@@ -6748,6 +6771,7 @@ mod tests {
                 invocation: protocol::FeatureInvocationIdentity("builtin:test/prepare".into()),
                 argument: Some("path".into()),
             }),
+            request_id: Some("completion-nonce".into()),
         };
         let wire: runtime_api::CompletionRequest =
             runtime_contract_convert(request.clone()).unwrap();

crates/workspace-server/src/server.rs

@@ -6330,6 +6330,7 @@ fn worker_completions_request_from_api(
         kind: request.kind,
         prefix: request.prefix,
         context: request.context,
+        request_id: request.request_id,
     })
 }
 
@@ -6427,6 +6428,8 @@ fn worker_completions_result_to_api(
         worker_id: result.worker.worker_id,
         kind: result.kind,
         prefix: result.prefix,
+        context: result.context,
+        request_id: result.request_id,
         entries: result.entries,
         diagnostics: result.diagnostics.into_iter().map(Into::into).collect(),
     })
@@ -35577,6 +35580,33 @@ mod tests {
         }
     }
 
+    #[test]
+    fn feature_completion_api_conversions_preserve_nonce_and_context_for_aba_requests() {
+        for nonce in ["first", "second", "first"] {
+            let json = serde_json::json!({"kind":"feature_argument", "prefix":"資料/", "request_id":nonce,
+                "context":{"invocation":"builtin:test/prepare", "argument":"path"}});
+            let wire: server_api::RuntimeWorkerCompletionsRequest =
+                serde_json::from_value(json.clone()).unwrap();
+            let request = worker_completions_request_from_api(wire).unwrap();
+            assert_eq!(serde_json::to_value(&request).unwrap(), json);
+            let result = worker_completions_result_to_api(WorkerCompletionsResult {
+                worker: RuntimeWorkerRef::new("runtime", "worker"),
+                kind: request.kind,
+                prefix: request.prefix,
+                context: request.context,
+                request_id: request.request_id,
+                entries: Vec::new(),
+                diagnostics: Vec::new(),
+            })
+            .unwrap();
+            assert_eq!(result.request_id.as_deref(), Some(nonce));
+            assert_eq!(
+                serde_json::to_value(result.context).unwrap(),
+                json["context"]
+            );
+        }
+    }
+
     #[test]
     fn flow_or_generic_worker_state_change_is_not_ticket_completion_authority() {
         let change = || {