@@ -57,6 +57,10 @@ static AGENT_DETECTORS: &[AgentDetector] = &[
5757 name : "sourcegraph-cody" ,
5858 env_vars : & [ "SRC_CODY" ] ,
5959 } ,
60+ AgentDetector {
61+ name : "pi-dev" ,
62+ env_vars : & [ "PI_CODING_AGENT" ] ,
63+ } ,
6064 AgentDetector {
6165 name : "generic-agent" ,
6266 env_vars : & [ "AGENT" ] ,
@@ -151,29 +155,36 @@ pub fn get_with_command(command: Option<&str>) -> String {
151155#[ cfg( test) ]
152156mod tests {
153157 use super :: * ;
158+ use crate :: test_utils:: ENV_LOCK ;
159+
160+ fn clear_all_agent_vars ( ) {
161+ for det in AGENT_DETECTORS {
162+ for var in det. env_vars {
163+ std:: env:: remove_var ( var) ;
164+ }
165+ }
166+ std:: env:: remove_var ( "FORCE_AGENT_MODE" ) ;
167+ }
154168
155169 #[ test]
156170 fn test_is_env_truthy ( ) {
157- // These tests use env vars that shouldn't be set in normal environments
171+ let _guard = ENV_LOCK . blocking_lock ( ) ;
158172 std:: env:: set_var ( "__PUP_TEST_TRUE__" , "true" ) ;
159173 assert ! ( is_env_truthy( "__PUP_TEST_TRUE__" ) ) ;
160-
161174 std:: env:: set_var ( "__PUP_TEST_ONE__" , "1" ) ;
162175 assert ! ( is_env_truthy( "__PUP_TEST_ONE__" ) ) ;
163-
164176 std:: env:: set_var ( "__PUP_TEST_FALSE__" , "false" ) ;
165177 assert ! ( !is_env_truthy( "__PUP_TEST_FALSE__" ) ) ;
166-
167178 assert ! ( !is_env_truthy( "__PUP_TEST_NONEXISTENT__" ) ) ;
168-
169- // Clean up
170179 std:: env:: remove_var ( "__PUP_TEST_TRUE__" ) ;
171180 std:: env:: remove_var ( "__PUP_TEST_ONE__" ) ;
172181 std:: env:: remove_var ( "__PUP_TEST_FALSE__" ) ;
173182 }
174183
175184 #[ test]
176185 fn test_user_agent_format ( ) {
186+ let _guard = ENV_LOCK . blocking_lock ( ) ;
187+ clear_all_agent_vars ( ) ;
177188 let ua = get ( ) ;
178189 assert ! ( ua. starts_with( "pup/" ) ) ;
179190 assert ! ( ua. contains( "rust" ) ) ;
@@ -184,6 +195,8 @@ mod tests {
184195
185196 #[ test]
186197 fn test_user_agent_with_command ( ) {
198+ let _guard = ENV_LOCK . blocking_lock ( ) ;
199+ clear_all_agent_vars ( ) ;
187200 let ua = get_with_command ( Some ( "security-findings-analyze" ) ) ;
188201 assert ! ( ua. starts_with( "pup/" ) ) ;
189202 assert ! ( ua. contains( "; cmd security-findings-analyze)" ) ) ;
@@ -196,6 +209,8 @@ mod tests {
196209/// until the UA is fixed.
197210#[ test]
198211 fn test_user_agent_parses_for_smart_edge_telemetry ( ) {
212+ let _guard = ENV_LOCK . blocking_lock ( ) ;
213+ clear_all_agent_vars ( ) ;
199214 for ua in [
200215 get_with_command ( None ) ,
201216 get_with_command ( Some ( "monitors-list" ) ) ,
@@ -228,6 +243,8 @@ mod tests {
228243
229244 #[ test]
230245 fn test_user_agent_with_no_command ( ) {
246+ let _guard = ENV_LOCK . blocking_lock ( ) ;
247+ clear_all_agent_vars ( ) ;
231248 let ua = get_with_command ( None ) ;
232249 assert ! ( !ua. contains( "cmd " ) ) ;
233250 assert_eq ! ( ua, get( ) ) ;
@@ -241,21 +258,17 @@ mod tests {
241258
242259 #[ test]
243260 fn test_detect_agent_info_no_agent ( ) {
244- // Clear all agent env vars
245- for det in AGENT_DETECTORS {
246- for var in det. env_vars {
247- std:: env:: remove_var ( var) ;
248- }
249- }
250- std:: env:: remove_var ( "FORCE_AGENT_MODE" ) ;
251-
261+ let _guard = ENV_LOCK . blocking_lock ( ) ;
262+ clear_all_agent_vars ( ) ;
252263 let info = detect_agent_info ( ) ;
253264 assert ! ( !info. detected) ;
254265 assert ! ( info. name. is_empty( ) ) ;
255266 }
256267
257268 #[ test]
258269 fn test_detect_agent_info_claude_code ( ) {
270+ let _guard = ENV_LOCK . blocking_lock ( ) ;
271+ clear_all_agent_vars ( ) ;
259272 std:: env:: set_var ( "CLAUDE_CODE" , "1" ) ;
260273 let info = detect_agent_info ( ) ;
261274 assert ! ( info. detected) ;
@@ -265,9 +278,8 @@ mod tests {
265278
266279 #[ test]
267280 fn test_detect_agent_info_cursor ( ) {
268- // Clear higher-priority detectors
269- std:: env:: remove_var ( "CLAUDECODE" ) ;
270- std:: env:: remove_var ( "CLAUDE_CODE" ) ;
281+ let _guard = ENV_LOCK . blocking_lock ( ) ;
282+ clear_all_agent_vars ( ) ;
271283 std:: env:: set_var ( "CURSOR_AGENT" , "true" ) ;
272284 let info = detect_agent_info ( ) ;
273285 assert ! ( info. detected) ;
@@ -277,32 +289,33 @@ mod tests {
277289
278290 #[ test]
279291 fn test_is_agent_mode_force ( ) {
292+ let _guard = ENV_LOCK . blocking_lock ( ) ;
293+ clear_all_agent_vars ( ) ;
280294 std:: env:: set_var ( "FORCE_AGENT_MODE" , "1" ) ;
281295 assert ! ( is_agent_mode( ) ) ;
282296 std:: env:: remove_var ( "FORCE_AGENT_MODE" ) ;
283297 }
284298
285299 #[ test]
286300 fn test_is_agent_mode_via_detector ( ) {
287- std:: env:: remove_var ( "FORCE_AGENT_MODE" ) ;
301+ let _guard = ENV_LOCK . blocking_lock ( ) ;
302+ clear_all_agent_vars ( ) ;
288303 std:: env:: set_var ( "CLAUDE_CODE" , "true" ) ;
289304 assert ! ( is_agent_mode( ) ) ;
290305 std:: env:: remove_var ( "CLAUDE_CODE" ) ;
291306 }
292307
293308 #[ test]
294309 fn test_is_agent_mode_false ( ) {
295- std:: env:: remove_var ( "FORCE_AGENT_MODE" ) ;
296- for det in AGENT_DETECTORS {
297- for var in det. env_vars {
298- std:: env:: remove_var ( var) ;
299- }
300- }
310+ let _guard = ENV_LOCK . blocking_lock ( ) ;
311+ clear_all_agent_vars ( ) ;
301312 assert ! ( !is_agent_mode( ) ) ;
302313 }
303314
304315 #[ test]
305316 fn test_user_agent_with_detected_agent ( ) {
317+ let _guard = ENV_LOCK . blocking_lock ( ) ;
318+ clear_all_agent_vars ( ) ;
306319 std:: env:: set_var ( "CLAUDE_CODE" , "1" ) ;
307320 let ua = get ( ) ;
308321 assert ! (
@@ -314,11 +327,8 @@ mod tests {
314327
315328 #[ test]
316329 fn test_user_agent_without_agent ( ) {
317- for det in AGENT_DETECTORS {
318- for var in det. env_vars {
319- std:: env:: remove_var ( var) ;
320- }
321- }
330+ let _guard = ENV_LOCK . blocking_lock ( ) ;
331+ clear_all_agent_vars ( ) ;
322332 let ua = get ( ) ;
323333 assert ! (
324334 !ua. contains( "ai-agent" ) ,
@@ -327,13 +337,31 @@ mod tests {
327337 assert ! ( ua. ends_with( ')' ) ) ;
328338 }
329339
340+ #[ test]
341+ fn test_detect_agent_info_pi_dev ( ) {
342+ let _guard = ENV_LOCK . blocking_lock ( ) ;
343+ clear_all_agent_vars ( ) ;
344+ std:: env:: set_var ( "PI_CODING_AGENT" , "true" ) ;
345+ let info = detect_agent_info ( ) ;
346+ assert ! ( info. detected) ;
347+ assert_eq ! ( info. name, "pi-dev" ) ;
348+ std:: env:: remove_var ( "PI_CODING_AGENT" ) ;
349+ }
350+
351+ #[ test]
352+ fn test_detect_agent_info_pi_dev_falsy ( ) {
353+ let _guard = ENV_LOCK . blocking_lock ( ) ;
354+ clear_all_agent_vars ( ) ;
355+ std:: env:: set_var ( "PI_CODING_AGENT" , "false" ) ;
356+ let info = detect_agent_info ( ) ;
357+ assert ! ( !info. detected) ;
358+ std:: env:: remove_var ( "PI_CODING_AGENT" ) ;
359+ }
360+
330361 #[ test]
331362 fn test_detect_agent_info_generic_agent ( ) {
332- for det in AGENT_DETECTORS {
333- for var in det. env_vars {
334- std:: env:: remove_var ( var) ;
335- }
336- }
363+ let _guard = ENV_LOCK . blocking_lock ( ) ;
364+ clear_all_agent_vars ( ) ;
337365 std:: env:: set_var ( "AGENT" , "1" ) ;
338366 let info = detect_agent_info ( ) ;
339367 assert ! ( info. detected) ;
0 commit comments