diff --git a/agent/agentproc/api.go b/agent/agentproc/api.go index 681faa05d5..0116114039 100644 --- a/agent/agentproc/api.go +++ b/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, }) diff --git a/agent/agentproc/api_test.go b/agent/agentproc/api_test.go index 38566bc85f..0d2677f5be 100644 --- a/agent/agentproc/api_test.go +++ b/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() diff --git a/agent/agentproc/process.go b/agent/agentproc/process.go index c797bb633e..b29446f3dc 100644 --- a/agent/agentproc/process.go +++ b/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 } diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index 566a23af4b..3dcbb248c9 100644 --- a/coderd/chatd/chatd.go +++ b/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, diff --git a/coderd/chatd/chattool/execute.go b/coderd/chatd/chattool/execute.go index 0b5553f84c..d9d9f61549 100644 --- a/coderd/chatd/chattool/execute.go +++ b/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 }