public_git
/
tart-guest-agent
mirror of https://github.com/cirruslabs/tart-guest-agent.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add exec overrides 

Co-authored-by: fkorotkov-oai <fkorotkov@openai.com>
This commit is contained in:
kroo-oai 2026-05-04 13:55:13 -04:00 committed by Fedor Korotkov
parent 30fe76ca7c
commit a0e561b6b2
3 changed files with 137 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

84
internal/rpc/agent.pb.go
View File

@ -1,7 +1,7 @@
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.36.6
// protoc (unknown)
// protoc-gen-go v1.36.11
// 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
3, // 1: ExecRequest.standard_input:type_name -> IOChunk
2, // 2: ExecRequest.terminal_resize:type_name -> TerminalSize
7, // 3: ExecResponse.exit:type_name -> ExecResponse.Exit
3, // 4: ExecResponse.standard_output:type_name -> IOChunk
3, // 5: ExecResponse.standard_error:type_name -> IOChunk
2, // 6: ExecRequest.Command.terminal_size:type_name -> TerminalSize
0, // 7: Agent.Exec:input_type -> ExecRequest
4, // 8: Agent.ResolveIP:input_type -> ResolveIPRequest
1, // 9: Agent.Exec:output_type -> ExecResponse
5, // 10: Agent.ResolveIP:output_type -> ResolveIPResponse
9, // [9:11] is the sub-list for method output_type
7, // [7:9] is the sub-list for method input_type
7, // [7:7] is the sub-list for extension type_name
7, // [7:7] is the sub-list for extension extendee
0, // [0:7] is the sub-list for field type_name
6, // 0: ExecRequest.command:type_name -> ExecRequest.Command
3, // 1: ExecRequest.standard_input:type_name -> IOChunk
2, // 2: ExecRequest.terminal_resize:type_name -> TerminalSize
8, // 3: ExecResponse.exit:type_name -> ExecResponse.Exit
3, // 4: ExecResponse.standard_output:type_name -> IOChunk
3, // 5: ExecResponse.standard_error:type_name -> IOChunk
2, // 6: ExecRequest.Command.terminal_size:type_name -> TerminalSize
7, // 7: ExecRequest.Command.env:type_name -> ExecRequest.Command.EnvEntry
0, // 8: Agent.Exec:input_type -> ExecRequest
4, // 9: Agent.ResolveIP:input_type -> ResolveIPRequest
1, // 10: Agent.Exec:output_type -> ExecResponse
5, // 11: Agent.ResolveIP:output_type -> ResolveIPResponse
10, // [10:12] is the sub-list for method output_type
8, // [8:10] is the sub-list for method input_type
8, // [8:8] is the sub-list for extension 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,
},

77
internal/rpc/exec.go
View File

@ -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

3
proto/rpc/agent.proto
View File

@ -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 {