返回 CodeWhale
extension_hooks.rs
根目录 / crates / tui / src / core / engine / tests / extension_hooks.rs
1 //! Real programmable host listeners on the actual Engine admission path.
2 use super::*;
3 use crate::llm_client::mock::{MockLlmClient, canned};
4
5 #[tokio::test]
6 async fn typescript_mod_rewrites_denies_and_regates_native_tools_on_the_real_engine() {
7 let Some(node) = crate::extension_host::tests::node_for_tests("typescript_mod_real_engine")
8 else {
9 return;
10 };
11 let _policy = crate::plugins::activation::TestPolicyGuard::extension_host(true);
12 let fixture = crate::extension_host::tests::FixturePlugins::new(&["hook-policy"]).await;
13 for (name, text) in [
14 ("before.txt", "original"),
15 ("after.txt", "rewritten content"),
16 ("blocked.txt", "must not be returned"),
17 ("malformed.txt", "must not be returned"),
18 ("throw.txt", "must not be returned"),
19 ("rewrite-action.txt", "original read"),
20 ] {
21 fs::write(fixture.workspace().join(name), text).unwrap();
22 }
23 let manager = fixture.manager(node);
24 let warm = manager.attach(fixture.registry());
25 warm.sync().await.expect("real host activation");
26 let _manager = crate::extension_host::TestManagerGuard::install(Arc::clone(&manager));
27 let mock = Arc::new(MockLlmClient::new(vec![
28 canned::tool_call_turn("rewrite", "read", r#"{"path":"before.txt"}"#),
29 canned::tool_call_turn("deny", "read", r#"{"path":"blocked.txt"}"#),
30 canned::tool_call_turn("malformed", "read", r#"{"path":"malformed.txt"}"#),
31 canned::tool_call_turn("throw", "read", r#"{"path":"throw.txt"}"#),
32 canned::tool_call_turn(
33 "action",
34 "File",
35 r#"{"action":"read","path":"rewrite-action.txt"}"#,
36 ),
37 canned::simple_text_turn("Finished checking the mod."),
38 canned::tool_call_turn("after-disable", "read", r#"{"path":"before.txt"}"#),
39 canned::simple_text_turn("The mod is disabled."),
40 ]));
41 let config = Config::default();
42 let mut engine_config = deterministic_engine_config(fixture.workspace());
43 engine_config.features.enable(Feature::ExtensionHost);
44 engine_config.plugin_registry = Some(fixture.registry());
45 let (engine, handle) = Engine::new_with_model_client(engine_config, &config, mock.clone());
46 manager.reconcile().await.unwrap();
47 let task = tokio::spawn(engine.run());
48 handle
49 .send(external_user_message_op(
50 "Check the programmable mod",
51 AppMode::Agent,
52 &config,
53 ))
54 .await
55 .unwrap();
56 let mut outcomes = HashMap::new();
57 let mut approvals = 0;
58 let mut rx = handle.rx_event.write().await;
59 loop {
60 match tokio::time::timeout(model_turn_event_timeout(), rx.recv())
61 .await
62 .unwrap()
63 .unwrap()
64 {
65 Event::ApprovalRequired { id, input, .. } => {
66 approvals += 1;
67 assert_eq!(
68 input,
69 serde_json::json!({
70 "action": "write", "path": "rewritten.txt", "content": "requires fresh approval"
71 }),
72 "the approval must be for the re-prepared revised write"
73 );
74 handle.deny_tool_call(id).await.unwrap();
75 }
76 Event::ToolCallComplete {
77 model_call: Some(call),
78 result,
79 ..
80 } => {
81 outcomes.insert(call.provider_id, result);
82 }
83 Event::TurnComplete { status, error, .. } => {
84 assert_eq!(status, TurnOutcomeStatus::Completed, "{error:?}");
85 break;
86 }
87 _ => {}
88 }
89 }
90 assert!(
91 outcomes
92 .remove("rewrite")
93 .unwrap()
94 .unwrap()
95 .content
96 .contains("rewritten content")
97 );
98 assert!(
99 outcomes
100 .remove("deny")
101 .unwrap()
102 .unwrap_err()
103 .to_string()
104 .contains("blocked by fixture")
105 );
106 for id in ["malformed", "throw"] {
107 assert!(
108 outcomes
109 .remove(id)
110 .unwrap()
111 .unwrap_err()
112 .to_string()
113 .contains("returned no verdict")
114 );
115 }
116 assert!(outcomes.remove("action").unwrap().is_err());
117 assert_eq!(approvals, 1, "only the revised write needs an approval");
118 assert!(!fixture.workspace().join("rewritten.txt").exists());
119 let first_request = mock.captured_requests().into_iter().next().unwrap();
120 let first_prompt = first_request
121 .messages
122 .iter()
123 .find_map(crate::runtime_handoff::extension_prompt_contributions_display)
124 .expect("the authored prompt contribution is delivered as a complete runtime snapshot");
125 assert!(first_prompt.contains("Use the repository style guide when preparing release notes."));
126 assert!(first_prompt.contains("repo-style"));
127 drop(rx);
128 let disabled = fixture.disable("hook-policy");
129 crate::extension_host::plugins_changed(disabled);
130 manager.reconcile().await.unwrap();
131 handle
132 .send(external_user_message_op(
133 "Read again after disabling",
134 AppMode::Agent,
135 &config,
136 ))
137 .await
138 .unwrap();
139 let mut rx = handle.rx_event.write().await;
140 let mut original = false;
141 loop {
142 match tokio::time::timeout(model_turn_event_timeout(), rx.recv())
143 .await
144 .unwrap()
145 .unwrap()
146 {
147 Event::ToolCallComplete {
148 model_call: Some(call),
149 result: Ok(result),
150 ..
151 } if call.provider_id == "after-disable" => {
152 original = result.content.contains("original")
153 && !result.content.contains("rewritten content");
154 }
155 Event::ApprovalRequired { .. } => panic!("disabled hook cannot request approval"),
156 Event::TurnComplete { status, error, .. } => {
157 assert_eq!(status, TurnOutcomeStatus::Completed, "{error:?}");
158 break;
159 }
160 _ => {}
161 }
162 }
163 assert!(original, "disabled listener must not rewrite the next turn");
164 let after_disable = mock.last_request().unwrap();
165 let latest_context = after_disable
166 .messages
167 .iter()
168 .rev()
169 .find_map(crate::runtime_handoff::extension_prompt_contributions_display)
170 .expect("disabling the authored section records a withdrawn runtime snapshot");
171 assert!(latest_context.contains("All earlier extension prompt contributions are withdrawn"));
172 assert!(
173 !latest_context.contains("Use the repository style guide when preparing release notes.")
174 );
175 assert!(!latest_context.contains("repo-style"));
176 drop(rx);
177 handle.send(Op::Shutdown).await.unwrap();
178 task.await.unwrap();
179 manager.shutdown().await;
180 }
181
181 lines RUST