Add exec overrides

Co-authored-by: fkorotkov-oai <fkorotkov@openai.com>

internal/rpc/agent.pb.go

@@ -1,7 +1,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
-// 	protoc-gen-go v1.36.6
+// 	protoc-gen-go v1.36.11
-// 	protoc        (unknown)
+// 	protoc        v7.34.1
 // source: rpc/agent.proto
 
 package rpc
@@ -401,6 +401,9 @@ type ExecRequest_Command struct {
 	Interactive   bool                   `protobuf:"varint,3,opt,name=interactive,proto3" json:"interactive,omitempty"`
 	Tty           bool                   `protobuf:"varint,4,opt,name=tty,proto3" json:"tty,omitempty"`
 	TerminalSize  *TerminalSize          `protobuf:"bytes,5,opt,name=terminal_size,json=terminalSize,proto3" json:"terminal_size,omitempty"`
+	Detach        bool                   `protobuf:"varint,6,opt,name=detach,proto3" json:"detach,omitempty"`
+	Env           map[string]string      `protobuf:"bytes,7,rep,name=env,proto3" json:"env,omitempty" protobuf_key:"bytes,1,opt,name=key" protobuf_val:"bytes,2,opt,name=value"`
+	Workdir       string                 `protobuf:"bytes,8,opt,name=workdir,proto3" json:"workdir,omitempty"`
 	unknownFields protoimpl.UnknownFields
 	sizeCache     protoimpl.SizeCache
 }
@@ -470,6 +473,27 @@ func (x *ExecRequest_Command) GetTerminalSize() *TerminalSize {
 	return nil
 }
 
+func (x *ExecRequest_Command) GetDetach() bool {
+	if x != nil {
+		return x.Detach
+	}
+	return false
+}
+
+func (x *ExecRequest_Command) GetEnv() map[string]string {
+	if x != nil {
+		return x.Env
+	}
+	return nil
+}
+
+func (x *ExecRequest_Command) GetWorkdir() string {
+	if x != nil {
+		return x.Workdir
+	}
+	return ""
+}
+
 type ExecResponse_Exit struct {
 	state         protoimpl.MessageState `protogen:"open.v1"`
 	Code          int32                  `protobuf:"varint,1,opt,name=code,proto3" json:"code,omitempty"`
@@ -479,7 +503,7 @@ type ExecResponse_Exit struct {
 
 func (x *ExecResponse_Exit) Reset() {
 	*x = ExecResponse_Exit{}
-	mi := &file_rpc_agent_proto_msgTypes[7]
+	mi := &file_rpc_agent_proto_msgTypes[8]
 	ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 	ms.StoreMessageInfo(mi)
 }
@@ -491,7 +515,7 @@ func (x *ExecResponse_Exit) String() string {
 func (*ExecResponse_Exit) ProtoMessage() {}
 
 func (x *ExecResponse_Exit) ProtoReflect() protoreflect.Message {
-	mi := &file_rpc_agent_proto_msgTypes[7]
+	mi := &file_rpc_agent_proto_msgTypes[8]
 	if x != nil {
 		ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
 		if ms.LoadMessageInfo() == nil {
@@ -518,17 +542,23 @@ var File_rpc_agent_proto protoreflect.FileDescriptor
 
 const file_rpc_agent_proto_rawDesc = "" +
 	"\n" +
-	"\x0frpc/agent.proto\x1a\x1bgoogle/protobuf/empty.proto\"\xd0\x02\n" +
+	"\x0frpc/agent.proto\x1a\x1bgoogle/protobuf/empty.proto\"\xeb\x03\n" +
 	"\vExecRequest\x120\n" +
 	"\acommand\x18\x01 \x01(\v2\x14.ExecRequest.CommandH\x00R\acommand\x121\n" +
 	"\x0estandard_input\x18\x02 \x01(\v2\b.IOChunkH\x00R\rstandardInput\x128\n" +
-	"\x0fterminal_resize\x18\x03 \x01(\v2\r.TerminalSizeH\x00R\x0eterminalResize\x1a\x99\x01\n" +
+	"\x0fterminal_resize\x18\x03 \x01(\v2\r.TerminalSizeH\x00R\x0eterminalResize\x1a\xb4\x02\n" +
 	"\aCommand\x12\x12\n" +
 	"\x04name\x18\x01 \x01(\tR\x04name\x12\x12\n" +
 	"\x04args\x18\x02 \x03(\tR\x04args\x12 \n" +
 	"\vinteractive\x18\x03 \x01(\bR\vinteractive\x12\x10\n" +
 	"\x03tty\x18\x04 \x01(\bR\x03tty\x122\n" +
-	"\rterminal_size\x18\x05 \x01(\v2\r.TerminalSizeR\fterminalSizeB\x06\n" +
+	"\rterminal_size\x18\x05 \x01(\v2\r.TerminalSizeR\fterminalSize\x12\x16\n" +
+	"\x06detach\x18\x06 \x01(\bR\x06detach\x12/\n" +
+	"\x03env\x18\a \x03(\v2\x1d.ExecRequest.Command.EnvEntryR\x03env\x12\x18\n" +
+	"\aworkdir\x18\b \x01(\tR\aworkdir\x1a6\n" +
+	"\bEnvEntry\x12\x10\n" +
+	"\x03key\x18\x01 \x01(\tR\x03key\x12\x14\n" +
+	"\x05value\x18\x02 \x01(\tR\x05value:\x028\x01B\x06\n" +
 	"\x04type\"\xc4\x01\n" +
 	"\fExecResponse\x12(\n" +
 	"\x04exit\x18\x01 \x01(\v2\x12.ExecResponse.ExitH\x00R\x04exit\x123\n" +
@@ -561,7 +591,7 @@ func file_rpc_agent_proto_rawDescGZIP() []byte {
 	return file_rpc_agent_proto_rawDescData
 }
 
-var file_rpc_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 8)
+var file_rpc_agent_proto_msgTypes = make([]protoimpl.MessageInfo, 9)
 var file_rpc_agent_proto_goTypes = []any{
 	(*ExecRequest)(nil),         // 0: ExecRequest
 	(*ExecResponse)(nil),        // 1: ExecResponse
@@ -570,25 +600,27 @@ var file_rpc_agent_proto_goTypes = []any{
 	(*ResolveIPRequest)(nil),    // 4: ResolveIPRequest
 	(*ResolveIPResponse)(nil),   // 5: ResolveIPResponse
 	(*ExecRequest_Command)(nil), // 6: ExecRequest.Command
-	(*ExecResponse_Exit)(nil),   // 7: ExecResponse.Exit
+	nil,                         // 7: ExecRequest.Command.EnvEntry
+	(*ExecResponse_Exit)(nil),   // 8: ExecResponse.Exit
 }
 var file_rpc_agent_proto_depIdxs = []int32{
-	6, // 0: ExecRequest.command:type_name -> ExecRequest.Command
+	6,  // 0: ExecRequest.command:type_name -> ExecRequest.Command
-	3, // 1: ExecRequest.standard_input:type_name -> IOChunk
+	3,  // 1: ExecRequest.standard_input:type_name -> IOChunk
-	2, // 2: ExecRequest.terminal_resize:type_name -> TerminalSize
+	2,  // 2: ExecRequest.terminal_resize:type_name -> TerminalSize
-	7, // 3: ExecResponse.exit:type_name -> ExecResponse.Exit
+	8,  // 3: ExecResponse.exit:type_name -> ExecResponse.Exit
-	3, // 4: ExecResponse.standard_output:type_name -> IOChunk
+	3,  // 4: ExecResponse.standard_output:type_name -> IOChunk
-	3, // 5: ExecResponse.standard_error:type_name -> IOChunk
+	3,  // 5: ExecResponse.standard_error:type_name -> IOChunk
-	2, // 6: ExecRequest.Command.terminal_size:type_name -> TerminalSize
+	2,  // 6: ExecRequest.Command.terminal_size:type_name -> TerminalSize
-	0, // 7: Agent.Exec:input_type -> ExecRequest
+	7,  // 7: ExecRequest.Command.env:type_name -> ExecRequest.Command.EnvEntry
-	4, // 8: Agent.ResolveIP:input_type -> ResolveIPRequest
+	0,  // 8: Agent.Exec:input_type -> ExecRequest
-	1, // 9: Agent.Exec:output_type -> ExecResponse
+	4,  // 9: Agent.ResolveIP:input_type -> ResolveIPRequest
-	5, // 10: Agent.ResolveIP:output_type -> ResolveIPResponse
+	1,  // 10: Agent.Exec:output_type -> ExecResponse
-	9, // [9:11] is the sub-list for method output_type
+	5,  // 11: Agent.ResolveIP:output_type -> ResolveIPResponse
-	7, // [7:9] is the sub-list for method input_type
+	10, // [10:12] is the sub-list for method output_type
-	7, // [7:7] is the sub-list for extension type_name
+	8,  // [8:10] is the sub-list for method input_type
-	7, // [7:7] is the sub-list for extension extendee
+	8,  // [8:8] is the sub-list for extension type_name
-	0, // [0:7] is the sub-list for field type_name
+	8,  // [8:8] is the sub-list for extension extendee
+	0,  // [0:8] is the sub-list for field type_name
 }
 
 func init() { file_rpc_agent_proto_init() }
@@ -612,7 +644,7 @@ func file_rpc_agent_proto_init() {
 			GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
 			RawDescriptor: unsafe.Slice(unsafe.StringData(file_rpc_agent_proto_rawDesc), len(file_rpc_agent_proto_rawDesc)),
 			NumEnums:      0,
-			NumMessages:   8,
+			NumMessages:   9,
 			NumExtensions: 0,
 			NumServices:   1,
 		},

internal/rpc/exec.go

@@ -14,6 +14,7 @@ import (
 	"os/exec"
 	"slices"
 	"strings"
+	"syscall"
 )
 
 const (
@@ -36,9 +37,47 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse])
 	zap.S().Infof("executing %s", formatCommandAndArgs(firstExecRequestCommand.Command.Name,
 		firstExecRequestCommand.Command.Args))
 
+	if firstExecRequestCommand.Command.Detach &&
+		(firstExecRequestCommand.Command.Interactive || firstExecRequestCommand.Command.Tty) {
+		return fmt.Errorf("detach cannot be used with interactive or tty")
+	}
+
 	// Execute the command
-	cmd := exec.CommandContext(stream.Context(), firstExecRequestCommand.Command.Name,
+	execCtx := stream.Context()
+	if firstExecRequestCommand.Command.Detach {
+		execCtx = context.Background()
+	}
+
+	cmd := exec.CommandContext(execCtx, firstExecRequestCommand.Command.Name,
 		firstExecRequestCommand.Command.Args...)
+	applyExecOverrides(cmd, firstExecRequestCommand.Command)
+
+	if firstExecRequestCommand.Command.Detach {
+		cmd.Stdout = io.Discard
+		cmd.Stderr = io.Discard
+		cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true}
+
+		if err := cmd.Start(); err != nil {
+			return err
+		}
+		if cmd.Process != nil {
+			if err := cmd.Process.Release(); err != nil {
+				return err
+			}
+		}
+
+		if err := stream.Send(&ExecResponse{
+			Type: &ExecResponse_Exit_{
+				Exit: &ExecResponse_Exit{
+					Code: 0,
+				},
+			},
+		}); err != nil && !errors.Is(err, context.Canceled) {
+			return err
+		}
+
+		return nil
+	}
 
 	var stdin io.WriteCloser
 	var stdout, stderr io.ReadCloser
@@ -239,6 +278,42 @@ func (rpc *RPC) Exec(stream grpc.BidiStreamingServer[ExecRequest, ExecResponse])
 	})
 }
 
+func applyExecOverrides(cmd *exec.Cmd, command *ExecRequest_Command) {
+	if command.Workdir != "" {
+		cmd.Dir = command.Workdir
+	}
+
+	if len(command.Env) > 0 {
+		cmd.Env = mergeEnv(command.Env)
+	}
+}
+
+func mergeEnv(overrides map[string]string) []string {
+	if len(overrides) == 0 {
+		return os.Environ()
+	}
+
+	envMap := make(map[string]string, len(overrides))
+	for _, entry := range os.Environ() {
+		parts := strings.SplitN(entry, "=", 2)
+		if len(parts) != 2 {
+			continue
+		}
+		envMap[parts[0]] = parts[1]
+	}
+
+	for key, value := range overrides {
+		envMap[key] = value
+	}
+
+	merged := make([]string, 0, len(envMap))
+	for key, value := range envMap {
+		merged = append(merged, key+"="+value)
+	}
+
+	return merged
+}
+
 func formatCommandAndArgs(name string, args []string) string {
 	var all []string
 

proto/rpc/agent.proto

@@ -16,6 +16,9 @@ message ExecRequest {
     bool interactive = 3;
     bool tty = 4;
     TerminalSize terminal_size = 5;
+    bool detach = 6;
+    map<string, string> env = 7;
+    string workdir = 8;
   }
 
   oneof type {