Implement exec_service gRPC server in Go

Author: afq984

MODULE.bazel

@@ -65,3 +65,19 @@ git_override(
 bazel_dep(name = "rules_cc", version = "0.2.14")
 bazel_dep(name = "platforms", version = "1.0.0")
 bazel_dep(name = "protobuf", version = "33.4")
+
+# --- Go Toolchain --------------------------------------------------------
+
+bazel_dep(name = "rules_go", version = "0.60.0")
+bazel_dep(name = "gazelle", version = "0.48.0")
+
+go_sdk = use_extension("@rules_go//go:extensions.bzl", "go_sdk")
+go_sdk.download(version = "1.24.12")
+
+go_deps = use_extension("@gazelle//:extensions.bzl", "go_deps")
+go_deps.from_file(go_mod = "//:go.mod")
+use_repo(
+    go_deps,
+    "org_golang_google_grpc",
+    "org_golang_google_protobuf",
+)

MODULE.bazel.lock

@@ -12,7 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+load("@rules_go//proto:def.bzl", "go_proto_library")
+
 proto_library(
     name = "exec_service_proto",
     srcs = ["exec_service.proto"],
 )
+
+go_proto_library(
+    name = "exec_service_go_proto",
+    compilers = ["@rules_go//proto:go_grpc"],
+    importpath = "github.com/google/agent-shell-tools/exec_service/execservicepb",
+    proto = ":exec_service_proto",
+    visibility = ["//visibility:public"],
+)

exec_service/BUILD.bazel

@@ -16,6 +16,8 @@ syntax = "proto3";
 
 package exec_service;
 
+option go_package = "github.com/google/agent-shell-tools/exec_service/execservicepb";
+
 // A simple command execution service.
 service ExecService {
   // Executes a command and streams the output until completion.

exec_service/exec_service.proto

@@ -0,0 +1,37 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+load("@rules_go//go:def.bzl", "go_library", "go_test")
+
+go_library(
+    name = "server",
+    srcs = ["server.go"],
+    importpath = "github.com/google/agent-shell-tools/exec_service/server",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//exec_service:exec_service_go_proto",
+    ],
+)
+
+go_test(
+    name = "server_test",
+    srcs = ["server_test.go"],
+    deps = [
+        ":server",
+        "//exec_service:exec_service_go_proto",
+        "@org_golang_google_grpc//:grpc",
+        "@org_golang_google_grpc//credentials/insecure",
+        "@org_golang_google_grpc//test/bufconn",
+    ],
+)

exec_service/server/BUILD.bazel

@@ -0,0 +1,133 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package server implements the ExecService gRPC server.
+package server
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"os/exec"
+	"syscall"
+	"time"
+
+	pb "github.com/google/agent-shell-tools/exec_service/execservicepb"
+)
+
+// ExecServer implements the ExecService gRPC service.
+// It runs shell commands and streams their output.
+type ExecServer struct {
+	pb.UnimplementedExecServiceServer
+}
+
+// RunCommand executes a shell command and streams output events until the
+// command exits. The command_line is interpreted by sh -c. Stdout and stderr
+// are merged into the output stream.
+func (s *ExecServer) RunCommand(req *pb.StartCommandRequest, stream pb.ExecService_RunCommandServer) error {
+	cmd := exec.Command("sh", "-c", req.GetCommandLine())
+	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+
+	if wd := req.GetWorkingDir(); wd != "" {
+		cmd.Dir = wd
+	}
+
+	pr, pw, err := os.Pipe()
+	if err != nil {
+		return sendError(stream, fmt.Sprintf("pipe: %v", err))
+	}
+	defer pr.Close()
+	cmd.Stdout = pw
+	cmd.Stderr = pw
+
+	if err := cmd.Start(); err != nil {
+		pw.Close()
+		return sendError(stream, fmt.Sprintf("start: %v", err))
+	}
+	pw.Close()
+
+	killGroup := func() {
+		syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
+	}
+
+	// Wait for the command concurrently. When it exits, expire the pipe
+	// read deadline so the read loop drains buffered output without
+	// blocking on background children that inherited the pipe fds.
+	waitCh := make(chan error, 1)
+	go func() {
+		err := cmd.Wait()
+		pr.SetReadDeadline(time.Now())
+		waitCh <- err
+	}()
+
+	// Kill the process group if the client disconnects.
+	done := make(chan struct{})
+	defer close(done)
+	go func() {
+		select {
+		case <-stream.Context().Done():
+			killGroup()
+		case <-done:
+		}
+	}()
+
+	buf := make([]byte, 4096)
+	for {
+		n, err := pr.Read(buf)
+		if n > 0 {
+			if sendErr := stream.Send(&pb.ServerEvent{
+				Event: &pb.ServerEvent_Output{Output: append([]byte(nil), buf[:n]...)},
+			}); sendErr != nil {
+				killGroup()
+				<-waitCh
+				return sendErr
+			}
+		}
+		if err != nil {
+			break
+		}
+	}
+
+	exitCode := int32(0)
+	var errMsg string
+	if err := <-waitCh; err != nil {
+		var exitErr *exec.ExitError
+		if errors.As(err, &exitErr) {
+			exitCode = int32(exitErr.ExitCode())
+		} else {
+			exitCode = -1
+			errMsg = err.Error()
+		}
+	}
+
+	return stream.Send(&pb.ServerEvent{
+		Event: &pb.ServerEvent_Exited{
+			Exited: &pb.ExitInfo{
+				ExitCode:     exitCode,
+				ErrorMessage: errMsg,
+			},
+		},
+	})
+}
+
+func sendError(stream pb.ExecService_RunCommandServer, msg string) error {
+	return stream.Send(&pb.ServerEvent{
+		Event: &pb.ServerEvent_Exited{
+			Exited: &pb.ExitInfo{
+				ExitCode:     -1,
+				ErrorMessage: msg,
+			},
+		},
+	})
+}