fix(agent): reap exited processes and scope process list by chat ID (#22944)

agent/agentproc/api.go

@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"sort"
 
 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
@@ -69,7 +70,12 @@ func (api *API) handleStartProcess(rw http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	proc, err := api.manager.start(req)
+	var chatID string
+	if id, _, ok := agentgit.ExtractChatContext(r); ok {
+		chatID = id.String()
+	}
+
+	proc, err := api.manager.start(req, chatID)
 	if err != nil {
 		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
 			Message: "Failed to start process.",
@@ -105,7 +111,28 @@ func (api *API) handleStartProcess(rw http.ResponseWriter, r *http.Request) {
 func (api *API) handleListProcesses(rw http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 
-	infos := api.manager.list()
+	var chatID string
+	if id, _, ok := agentgit.ExtractChatContext(r); ok {
+		chatID = id.String()
+	}
+
+	infos := api.manager.list(chatID)
+
+	// Sort by running state (running first), then by started_at
+	// descending so the most recent processes appear first.
+	sort.Slice(infos, func(i, j int) bool {
+		if infos[i].Running != infos[j].Running {
+			return infos[i].Running
+		}
+		return infos[i].StartedAt > infos[j].StartedAt
+	})
+
+	// Cap the response to avoid bloating LLM context.
+	const maxListProcesses = 10
+	if len(infos) > maxListProcesses {
+		infos = infos[:maxListProcesses]
+	}
+
 	httpapi.Write(ctx, rw, http.StatusOK, workspacesdk.ListProcessesResponse{
 		Processes: infos,
 	})

agent/agentproc/api_test.go

@@ -27,7 +27,7 @@ import (
 )
 
 // postStart sends a POST /start request and returns the recorder.
-func postStart(t *testing.T, handler http.Handler, req workspacesdk.StartProcessRequest) *httptest.ResponseRecorder {
+func postStart(t *testing.T, handler http.Handler, req workspacesdk.StartProcessRequest, headers ...http.Header) *httptest.ResponseRecorder {
 	t.Helper()
 
 	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
@@ -38,6 +38,13 @@ func postStart(t *testing.T, handler http.Handler, req workspacesdk.StartProcess
 
 	w := httptest.NewRecorder()
 	r := httptest.NewRequestWithContext(ctx, http.MethodPost, "/start", bytes.NewReader(body))
+	for _, h := range headers {
+		for k, vals := range h {
+			for _, v := range vals {
+				r.Header.Add(k, v)
+			}
+		}
+	}
 	handler.ServeHTTP(w, r)
 	return w
 }
@@ -140,10 +147,10 @@ func waitForExit(t *testing.T, handler http.Handler, id string) workspacesdk.Pro
 
 // startAndGetID is a helper that starts a process and returns
 // the process ID.
-func startAndGetID(t *testing.T, handler http.Handler, req workspacesdk.StartProcessRequest) string {
+func startAndGetID(t *testing.T, handler http.Handler, req workspacesdk.StartProcessRequest, headers ...http.Header) string {
 	t.Helper()
 
-	w := postStart(t, handler, req)
+	w := postStart(t, handler, req, headers...)
 	require.Equal(t, http.StatusOK, w.Code)
 
 	var resp workspacesdk.StartProcessResponse
@@ -333,6 +340,180 @@ func TestListProcesses(t *testing.T) {
 		require.Empty(t, resp.Processes)
 	})
 
+	t.Run("FilterByChatID", func(t *testing.T) {
+		t.Parallel()
+
+		handler := newTestAPI(t)
+
+		chatA := uuid.New().String()
+		chatB := uuid.New().String()
+		headersA := http.Header{workspacesdk.CoderChatIDHeader: {chatA}}
+		headersB := http.Header{workspacesdk.CoderChatIDHeader: {chatB}}
+
+		// Start processes with different chat IDs.
+		id1 := startAndGetID(t, handler, workspacesdk.StartProcessRequest{
+			Command: "echo chat-a",
+		}, headersA)
+		waitForExit(t, handler, id1)
+
+		id2 := startAndGetID(t, handler, workspacesdk.StartProcessRequest{
+			Command: "echo chat-b",
+		}, headersB)
+		waitForExit(t, handler, id2)
+
+		id3 := startAndGetID(t, handler, workspacesdk.StartProcessRequest{
+			Command: "echo chat-a-2",
+		}, headersA)
+		waitForExit(t, handler, id3)
+
+		// List with chat A header should return 2 processes.
+		w := getListWithChatHeader(t, handler, chatA)
+		require.Equal(t, http.StatusOK, w.Code)
+
+		var resp workspacesdk.ListProcessesResponse
+		err := json.NewDecoder(w.Body).Decode(&resp)
+		require.NoError(t, err)
+		require.Len(t, resp.Processes, 2)
+
+		ids := make(map[string]bool)
+		for _, p := range resp.Processes {
+			ids[p.ID] = true
+		}
+		require.True(t, ids[id1])
+		require.True(t, ids[id3])
+
+		// List with chat B header should return 1 process.
+		w2 := getListWithChatHeader(t, handler, chatB)
+		require.Equal(t, http.StatusOK, w2.Code)
+
+		var resp2 workspacesdk.ListProcessesResponse
+		err = json.NewDecoder(w2.Body).Decode(&resp2)
+		require.NoError(t, err)
+		require.Len(t, resp2.Processes, 1)
+		require.Equal(t, id2, resp2.Processes[0].ID)
+
+		// List without chat header should return all 3.
+		w3 := getList(t, handler)
+		require.Equal(t, http.StatusOK, w3.Code)
+
+		var resp3 workspacesdk.ListProcessesResponse
+		err = json.NewDecoder(w3.Body).Decode(&resp3)
+		require.NoError(t, err)
+		require.Len(t, resp3.Processes, 3)
+	})
+
+	t.Run("ChatIDFiltering", func(t *testing.T) {
+		t.Parallel()
+
+		handler := newTestAPI(t)
+		chatID := uuid.New().String()
+		headers := http.Header{workspacesdk.CoderChatIDHeader: {chatID}}
+
+		id := startAndGetID(t, handler, workspacesdk.StartProcessRequest{
+			Command: "echo with-chat",
+		}, headers)
+		waitForExit(t, handler, id)
+
+		// Listing with the same chat header should return
+		// the process.
+		w := getListWithChatHeader(t, handler, chatID)
+		require.Equal(t, http.StatusOK, w.Code)
+
+		var resp workspacesdk.ListProcessesResponse
+		err := json.NewDecoder(w.Body).Decode(&resp)
+		require.NoError(t, err)
+		require.Len(t, resp.Processes, 1)
+		require.Equal(t, id, resp.Processes[0].ID)
+
+		// Listing with a different chat header should not
+		// return the process.
+		w2 := getListWithChatHeader(t, handler, uuid.New().String())
+		require.Equal(t, http.StatusOK, w2.Code)
+
+		var resp2 workspacesdk.ListProcessesResponse
+		err = json.NewDecoder(w2.Body).Decode(&resp2)
+		require.NoError(t, err)
+		require.Empty(t, resp2.Processes)
+
+		// Listing without a chat header should return the
+		// process (no filtering).
+		w3 := getList(t, handler)
+		require.Equal(t, http.StatusOK, w3.Code)
+
+		var resp3 workspacesdk.ListProcessesResponse
+		err = json.NewDecoder(w3.Body).Decode(&resp3)
+		require.NoError(t, err)
+		require.Len(t, resp3.Processes, 1)
+	})
+
+	t.Run("SortAndLimit", func(t *testing.T) {
+		t.Parallel()
+
+		handler := newTestAPI(t)
+
+		// Start 12 short-lived processes so we exceed the
+		// limit of 10.
+		for i := 0; i < 12; i++ {
+			id := startAndGetID(t, handler, workspacesdk.StartProcessRequest{
+				Command: fmt.Sprintf("echo proc-%d", i),
+			})
+			waitForExit(t, handler, id)
+		}
+
+		w := getList(t, handler)
+		require.Equal(t, http.StatusOK, w.Code)
+
+		var resp workspacesdk.ListProcessesResponse
+		err := json.NewDecoder(w.Body).Decode(&resp)
+		require.NoError(t, err)
+		require.Len(t, resp.Processes, 10, "should be capped at 10")
+
+		// All returned processes are exited, so they should
+		// be sorted by StartedAt descending (newest first).
+		for i := 1; i < len(resp.Processes); i++ {
+			require.GreaterOrEqual(t, resp.Processes[i-1].StartedAt, resp.Processes[i].StartedAt,
+				"processes should be sorted by started_at descending")
+		}
+	})
+
+	t.Run("RunningProcessesSortedFirst", func(t *testing.T) {
+		t.Parallel()
+
+		handler := newTestAPI(t)
+
+		// Start an exited process first.
+		exitedID := startAndGetID(t, handler, workspacesdk.StartProcessRequest{
+			Command: "echo done",
+		})
+		waitForExit(t, handler, exitedID)
+
+		// Start a running process after.
+		runningID := startAndGetID(t, handler, workspacesdk.StartProcessRequest{
+			Command:    "sleep 300",
+			Background: true,
+		})
+
+		w := getList(t, handler)
+		require.Equal(t, http.StatusOK, w.Code)
+
+		var resp workspacesdk.ListProcessesResponse
+		err := json.NewDecoder(w.Body).Decode(&resp)
+		require.NoError(t, err)
+		require.Len(t, resp.Processes, 2)
+
+		// Running process should come first regardless of
+		// start order.
+		require.Equal(t, runningID, resp.Processes[0].ID)
+		require.True(t, resp.Processes[0].Running)
+		require.Equal(t, exitedID, resp.Processes[1].ID)
+		require.False(t, resp.Processes[1].Running)
+
+		// Clean up.
+		postSignal(t, handler, runningID, workspacesdk.SignalProcessRequest{
+			Signal: "kill",
+		})
+	})
+
 	t.Run("MixedRunningAndExited", func(t *testing.T) {
 		t.Parallel()
 
@@ -381,6 +562,23 @@ func TestListProcesses(t *testing.T) {
 	})
 }
 
+// getListWithChatHeader sends a GET /list request with the
+// Coder-Chat-Id header set and returns the recorder.
+func getListWithChatHeader(t *testing.T, handler http.Handler, chatID string) *httptest.ResponseRecorder {
+	t.Helper()
+
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
+	defer cancel()
+
+	w := httptest.NewRecorder()
+	r := httptest.NewRequestWithContext(ctx, http.MethodGet, "/list", nil)
+	if chatID != "" {
+		r.Header.Set(workspacesdk.CoderChatIDHeader, chatID)
+	}
+	handler.ServeHTTP(w, r)
+	return w
+}
+
 func TestProcessOutput(t *testing.T) {
 	t.Parallel()
 

agent/agentproc/process.go

@@ -21,6 +21,10 @@ import (
 var (
 	errProcessNotFound   = xerrors.New("process not found")
 	errProcessNotRunning = xerrors.New("process is not running")
+
+	// exitedProcessReapAge is how long an exited process is
+	// kept before being automatically removed from the map.
+	exitedProcessReapAge = 5 * time.Minute
 )
 
 // process represents a running or completed process.
@@ -30,6 +34,7 @@ type process struct {
 	command    string
 	workDir    string
 	background bool
+	chatID     string
 	cmd        *exec.Cmd
 	cancel     context.CancelFunc
 	buf        *HeadTailBuffer
@@ -89,7 +94,7 @@ func newManager(logger slog.Logger, execer agentexec.Execer, updateEnv func(curr
 // processes use a long-lived context so the process survives
 // the HTTP request lifecycle. The background flag only affects
 // client-side polling behavior.
-func (m *manager) start(req workspacesdk.StartProcessRequest) (*process, error) {
+func (m *manager) start(req workspacesdk.StartProcessRequest, chatID string) (*process, error) {
 	m.mu.Lock()
 	if m.closed {
 		m.mu.Unlock()
@@ -154,6 +159,7 @@ func (m *manager) start(req workspacesdk.StartProcessRequest) (*process, error)
 		command:    req.Command,
 		workDir:    req.WorkDir,
 		background: req.Background,
+		chatID:     chatID,
 		cmd:        cmd,
 		cancel:     cancel,
 		buf:        buf,
@@ -215,14 +221,32 @@ func (m *manager) get(id string) (*process, bool) {
 	return proc, ok
 }
 
-// list returns info about all tracked processes.
-func (m *manager) list() []workspacesdk.ProcessInfo {
+// list returns info about all tracked processes. Exited
+// processes older than exitedProcessReapAge are removed.
+// If chatID is non-empty, only processes belonging to that
+// chat are returned.
+func (m *manager) list(chatID string) []workspacesdk.ProcessInfo {
 	m.mu.Lock()
 	defer m.mu.Unlock()
 
+	now := m.clock.Now()
 	infos := make([]workspacesdk.ProcessInfo, 0, len(m.procs))
-	for _, proc := range m.procs {
-		infos = append(infos, proc.info())
+	for id, proc := range m.procs {
+		info := proc.info()
+		// Reap processes that exited more than 5 minutes ago
+		// to prevent unbounded map growth.
+		if !info.Running && info.ExitedAt != nil {
+			exitedAt := time.Unix(*info.ExitedAt, 0)
+			if now.Sub(exitedAt) > exitedProcessReapAge {
+				delete(m.procs, id)
+				continue
+			}
+		}
+		// Filter by chatID if provided.
+		if chatID != "" && proc.chatID != chatID {
+			continue
+		}
+		infos = append(infos, info)
 	}
 	return infos
 }

coderd/chatd/chatd.go

@@ -2482,7 +2482,6 @@ func (p *Server) runChat(
 		}),
 		chattool.Execute(chattool.ExecuteOptions{
 			GetWorkspaceConn: getWorkspaceConn,
-			ChatID:           chat.ID.String(),
 		}),
 		chattool.ProcessOutput(chattool.ProcessToolOptions{
 			GetWorkspaceConn: getWorkspaceConn,

coderd/chatd/chattool/execute.go

@@ -65,7 +65,6 @@ type ExecuteResult struct {
 type ExecuteOptions struct {
 	GetWorkspaceConn func(context.Context) (workspacesdk.AgentConn, error)
 	DefaultTimeout   time.Duration
-	ChatID           string
 }
 
 // ProcessToolOptions configures a process management tool
@@ -97,7 +96,7 @@ func Execute(options ExecuteOptions) fantasy.AgentTool {
 			if err != nil {
 				return fantasy.NewTextErrorResponse(err.Error()), nil
 			}
-			return executeTool(ctx, conn, args, options.DefaultTimeout, options.ChatID), nil
+			return executeTool(ctx, conn, args, options.DefaultTimeout), nil
 		},
 	)
 }
@@ -107,7 +106,6 @@ func executeTool(
 	conn workspacesdk.AgentConn,
 	args ExecuteArgs,
 	optTimeout time.Duration,
-	chatID string,
 ) fantasy.ToolResponse {
 	if args.Command == "" {
 		return fantasy.NewTextErrorResponse("command is required")
@@ -116,9 +114,6 @@ func executeTool(
 	// Build the environment map for the process request.
 	env := make(map[string]string, len(nonInteractiveEnvVars)+1)
 	env["CODER_CHAT_AGENT"] = "true"
-	if chatID != "" {
-		env["CODER_CHAT_ID"] = chatID
-	}
 	for k, v := range nonInteractiveEnvVars {
 		env[k] = v
 	}