@@ -45,7 +45,10 @@ class TestResultMessage:
4545 async def test_structured_output_result (self , fake_transport , make_query_options ):
4646 result_msg = {
4747 "message_type" : "result_message" ,
48- "data" : {"type" : "structured_output" , "structured_output" : {"success" : True }},
48+ "data" : {
49+ "type" : "structured_output" ,
50+ "structured_output" : {"success" : True },
51+ },
4952 "metadata" : {"duration_ms" : 50 },
5053 }
5154 transport , client = fake_transport ([result_msg ])
@@ -363,7 +366,10 @@ async def test_full_message_sequence(self, fake_transport, make_query_options):
363366 },
364367 {
365368 "message_type" : "result_message" ,
366- "data" : {"type" : "structured_output" , "structured_output" : {"answer" : 42 }},
369+ "data" : {
370+ "type" : "structured_output" ,
371+ "structured_output" : {"answer" : 42 },
372+ },
367373 "metadata" : {"duration_ms" : 100 },
368374 },
369375 ]
@@ -421,6 +427,36 @@ async def test_schema_forwarded_when_set(self, fake_transport, make_query_option
421427 invocation_msgs [0 ]["output_format_json_schema_str" ] == '{"type":"object"}'
422428 )
423429
430+ @pytest .mark .asyncio
431+ async def test_agent_env_forwarded_when_set (
432+ self , fake_transport , make_query_options
433+ ):
434+ transport , client = fake_transport ([TEXT_RESULT_MSG ])
435+
436+ await client .query (
437+ prompt = DEFAULT_PROMPT ,
438+ options = make_query_options (agent_env = {"MY_VAR" : "hello" }),
439+ )
440+
441+ invocation_msgs = [
442+ m for m in transport .sent if m .get ("message_type" ) == "invocation_message"
443+ ]
444+ assert invocation_msgs [0 ]["agent_env" ] == {"MY_VAR" : "hello" }
445+
446+ @pytest .mark .asyncio
447+ async def test_agent_env_none_when_omitted (self , fake_transport , query_options ):
448+ transport , client = fake_transport ([TEXT_RESULT_MSG ])
449+
450+ await client .query (
451+ prompt = DEFAULT_PROMPT ,
452+ options = query_options ,
453+ )
454+
455+ invocation_msgs = [
456+ m for m in transport .sent if m .get ("message_type" ) == "invocation_message"
457+ ]
458+ assert invocation_msgs [0 ]["agent_env" ] is None
459+
424460 @pytest .mark .asyncio
425461 async def test_schema_none_when_omitted (self , fake_transport , query_options ):
426462 transport , client = fake_transport ([TEXT_RESULT_MSG ])
@@ -518,7 +554,9 @@ async def test_sends_command_execution_message(
518554 ]
519555 assert len (cmd_msgs ) == 1
520556 assert cmd_msgs [0 ]["source_id" ] == "cmd-test"
521- assert cmd_msgs [0 ]["commands" ] == [{"command" : "echo hello" , "cwd" : None }]
557+ assert cmd_msgs [0 ]["commands" ] == [
558+ {"command" : "echo hello" , "cwd" : None , "env" : None }
559+ ]
522560
523561 @pytest .mark .asyncio
524562 async def test_assistant_message_dispatched (
@@ -641,6 +679,49 @@ async def on_artifact(
641679 assert response_msgs [0 ]["filename" ] == "output.txt"
642680 assert response_msgs [0 ]["presigned_url" ] == "https://s3.example.com/presigned"
643681
682+ @pytest .mark .asyncio
683+ async def test_command_env_forwarded (
684+ self , fake_transport , make_execute_commands_options
685+ ):
686+ result_msg = {
687+ "message_type" : "command_execution_result_message" ,
688+ "results" : [{"command" : "echo hello" , "exit_code" : 0 }],
689+ }
690+ transport , client = fake_transport ([result_msg ])
691+
692+ await client .execute_commands (
693+ commands = [CommandInterface (command = "echo hello" , env = {"FOO" : "bar" })],
694+ options = make_execute_commands_options (),
695+ )
696+
697+ cmd_msgs = [
698+ m
699+ for m in transport .sent
700+ if m .get ("message_type" ) == "command_execution_message"
701+ ]
702+ assert len (cmd_msgs ) == 1
703+ assert cmd_msgs [0 ]["commands" ] == [
704+ {"command" : "echo hello" , "cwd" : None , "env" : {"FOO" : "bar" }}
705+ ]
706+
707+ @pytest .mark .asyncio
708+ async def test_command_env_none_by_default (
709+ self , fake_transport , make_execute_commands_options
710+ ):
711+ transport , client = fake_transport ([COMMAND_RESULT_MSG ])
712+
713+ await client .execute_commands (
714+ commands = [CommandInterface (command = "echo hello" )],
715+ options = make_execute_commands_options (),
716+ )
717+
718+ cmd_msgs = [
719+ m
720+ for m in transport .sent
721+ if m .get ("message_type" ) == "command_execution_message"
722+ ]
723+ assert cmd_msgs [0 ]["commands" ][0 ]["env" ] is None
724+
644725 def test_execute_commands_options_artifacts_validation (self ):
645726 with pytest .raises (ValueError , match = "must be specified together" ):
646727 ExecuteCommandsOptions (artifacts_dir = "/tmp/artifacts" )
0 commit comments