bkinnightskytw
/
agent-orchestrator
mirror of https://github.com/ComposioHQ/agent-orchestrator.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix(cli): verify daemon ownership before stop signal

This commit is contained in:
Dhruv Sharma 2026-05-31 21:45:00 +05:30 committed by harshitsinghbhandari
parent f72facb9e5
commit 925e70763d
5 changed files with 118 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
backend/internal/cli/root_test.go
View File

@ -2,6 +2,7 @@ package cli
import (
"bytes"
"fmt"
"net"
"net/http"
"net/http/httptest"
@ -13,6 +14,7 @@ import (
"testing"
"time"
"github.com/aoagents/agent-orchestrator/backend/internal/daemonmeta"
"github.com/aoagents/agent-orchestrator/backend/internal/runfile"
)
@ -51,9 +53,9 @@ func TestStartReturnsExistingReadyDaemon(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/healthz":
_, _ = w.Write([]byte(`{"status":"ok"}`))
_, _ = fmt.Fprintf(w, `{"status":"ok","service":%q,"pid":%d}`, daemonmeta.ServiceName, os.Getpid())
case "/readyz":
_, _ = w.Write([]byte(`{"status":"ready"}`))
_, _ = fmt.Fprintf(w, `{"status":"ready","service":%q,"pid":%d}`, daemonmeta.ServiceName, os.Getpid())
default:
http.NotFound(w, r)
}
@ -107,6 +109,50 @@ func TestStopRemovesStaleRunFile(t *testing.T) {
}
}
func TestStopDoesNotSignalUnverifiedReusedPID(t *testing.T) {
cfg := setConfigEnv(t)
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch r.URL.Path {
case "/healthz":
_, _ = w.Write([]byte(`{"status":"ok"}`))
case "/readyz":
_, _ = w.Write([]byte(`{"status":"ready"}`))
default:
http.NotFound(w, r)
}
}))
defer srv.Close()
if err := runfile.Write(cfg.runFile, runfile.Info{PID: 4242, Port: serverPort(t, srv.URL), StartedAt: time.Unix(100, 0).UTC()}); err != nil {
t.Fatal(err)
}
var signaled bool
out, _, err := executeCLI(t, Deps{
ProcessAlive: func(pid int) bool { return pid == 4242 },
SignalTerm: func(pid int) error {
signaled = true
return nil
},
}, "stop", "--json")
if err != nil {
t.Fatal(err)
}
if signaled {
t.Fatal("stop signaled a PID whose health probe did not prove AO daemon ownership")
}
if !strings.Contains(out, `"state": "stopped"`) {
t.Fatalf("stop did not report stopped:\n%s", out)
}
info, err := runfile.Read(cfg.runFile)
if err != nil {
t.Fatal(err)
}
if info != nil {
t.Fatalf("unverified run-file was not removed: %#v", info)
}
}
type testConfig struct {
runFile string
dataDir string

60
backend/internal/cli/status.go
View File

@ -10,6 +10,7 @@ import (
"github.com/spf13/cobra"
"github.com/aoagents/agent-orchestrator/backend/internal/config"
"github.com/aoagents/agent-orchestrator/backend/internal/daemonmeta"
"github.com/aoagents/agent-orchestrator/backend/internal/runfile"
)
@ -30,6 +31,13 @@ type daemonStatus struct {
Health string `json:"health,omitempty"`
Ready string `json:"ready,omitempty"`
Error string `json:"error,omitempty"`
owned bool
}
type probeResult struct {
Status string `json:"status"`
Service string `json:"service"`
PID int `json:"pid"`
}
func newStatusCommand(ctx *commandContext) *cobra.Command {
@ -81,11 +89,21 @@ func (c *commandContext) inspectDaemon(ctx context.Context) (daemonStatus, error
health, err := c.readProbe(ctx, info.Port, "healthz")
if err != nil {
st.State = "unhealthy"
st.State = "stale"
st.Error = err.Error()
return st, nil
}
st.Health = health
if err := verifyProbeOwner(health, info.PID, "healthz"); err != nil {
st.State = "stale"
st.Error = err.Error()
return st, nil
}
st.owned = true
st.Health = health.Status
if health.Status != "ok" {
st.State = "unhealthy"
return st, nil
}
ready, err := c.readProbe(ctx, info.Port, "readyz")
if err != nil {
@ -93,8 +111,14 @@ func (c *commandContext) inspectDaemon(ctx context.Context) (daemonStatus, error
st.Error = err.Error()
return st, nil
}
st.Ready = ready
if ready == "ready" {
if err := verifyProbeOwner(ready, info.PID, "readyz"); err != nil {
st.State = "stale"
st.owned = false
st.Error = err.Error()
return st, nil
}
st.Ready = ready.Status
if ready.Status == "ready" {
st.State = "ready"
return st, nil
}
@ -102,32 +126,40 @@ func (c *commandContext) inspectDaemon(ctx context.Context) (daemonStatus, error
return st, nil
}
func (c *commandContext) readProbe(ctx context.Context, port int, path string) (string, error) {
func (c *commandContext) readProbe(ctx context.Context, port int, path string) (probeResult, error) {
reqCtx, cancel := context.WithTimeout(ctx, probeTimeout)
defer cancel()
req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, fmt.Sprintf("http://%s:%d/%s", config.LoopbackHost, port, path), nil)
if err != nil {
return "", err
return probeResult{}, err
}
resp, err := c.deps.HTTPClient.Do(req)
if err != nil {
return "", fmt.Errorf("%s: %w", path, err)
return probeResult{}, fmt.Errorf("%s: %w", path, err)
}
defer resp.Body.Close()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
return "", fmt.Errorf("%s: HTTP %d", path, resp.StatusCode)
}
var body struct {
Status string `json:"status"`
return probeResult{}, fmt.Errorf("%s: HTTP %d", path, resp.StatusCode)
}
var body probeResult
if err := json.NewDecoder(resp.Body).Decode(&body); err != nil {
return "", fmt.Errorf("%s: decode response: %w", path, err)
return probeResult{}, fmt.Errorf("%s: decode response: %w", path, err)
}
if body.Status == "" {
return "", fmt.Errorf("%s: missing status", path)
return probeResult{}, fmt.Errorf("%s: missing status", path)
}
return body.Status, nil
return body, nil
}
func verifyProbeOwner(probe probeResult, wantPID int, path string) error {
if probe.Service != daemonmeta.ServiceName {
return fmt.Errorf("%s: response is not from AO daemon", path)
}
if probe.PID != wantPID {
return fmt.Errorf("%s: daemon pid %d does not match run-file pid %d", path, probe.PID, wantPID)
}
return nil
}
func writeStatus(cmd *cobra.Command, st daemonStatus) error {

6
backend/internal/cli/stop.go
View File

@ -61,6 +61,12 @@ func (c *commandContext) stopDaemon(ctx context.Context, opts stopOptions) (daem
}
return daemonStatus{State: "stopped", RunFile: cfg.RunFilePath, DataDir: cfg.DataDir}, nil
}
if !st.owned {
if err := runfile.Remove(cfg.RunFilePath); err != nil {
return daemonStatus{}, err
}
return daemonStatus{State: "stopped", RunFile: cfg.RunFilePath, DataDir: cfg.DataDir}, nil
}
if err := c.deps.SignalTerm(st.PID); err != nil {
if c.deps.ProcessAlive(st.PID) {

6
backend/internal/daemonmeta/meta.go Normal file
View File

@ -0,0 +1,6 @@
package daemonmeta
// ServiceName identifies the AO daemon in loopback health/readiness probes.
// The CLI uses it with the reported PID to avoid signaling an unrelated process
// when a stale run-file's PID has been reused.
const ServiceName = "agent-orchestrator-daemon"

14
backend/internal/httpd/router.go
View File

@ -7,11 +7,13 @@ package httpd
import (
"log/slog"
"net/http"
"os"
"github.com/go-chi/chi/v5"
"github.com/go-chi/chi/v5/middleware"
"github.com/aoagents/agent-orchestrator/backend/internal/config"
"github.com/aoagents/agent-orchestrator/backend/internal/daemonmeta"
"github.com/aoagents/agent-orchestrator/backend/internal/terminal"
)
@ -68,12 +70,20 @@ func mountHealth(r chi.Router) {
// handleHealthz is the liveness probe: it answers 200 as long as the process is
// up and serving. It does no dependency checks by design.
func handleHealthz(w http.ResponseWriter, _ *http.Request) {
writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
writeJSON(w, http.StatusOK, map[string]any{
"status": "ok",
"service": daemonmeta.ServiceName,
"pid": os.Getpid(),
})
}
// handleReadyz is the readiness probe. In the 1a skeleton the daemon is ready
// as soon as it is listening; later phases will gate this on dependency
// initialisation (e.g. store/event-bus warm-up).
func handleReadyz(w http.ResponseWriter, _ *http.Request) {
writeJSON(w, http.StatusOK, map[string]string{"status": "ready"})
writeJSON(w, http.StatusOK, map[string]any{
"status": "ready",
"service": daemonmeta.ServiceName,
"pid": os.Getpid(),
})
}