diff --git a/agent/agent.go b/agent/agent.go index d876493eab..8826426087 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -48,6 +48,7 @@ import ( "github.com/coder/coder/v2/agent/agentscripts" "github.com/coder/coder/v2/agent/agentsocket" "github.com/coder/coder/v2/agent/agentssh" + "github.com/coder/coder/v2/agent/agenttoolcall" "github.com/coder/coder/v2/agent/boundarylogproxy" "github.com/coder/coder/v2/agent/proto" "github.com/coder/coder/v2/agent/proto/resourcesmonitor" @@ -344,6 +345,7 @@ type agent struct { filesAPI *agentfiles.API gitAPI *agentgit.API processAPI *agentproc.API + toolCalls *agenttoolcall.Table desktopAPI *agentdesktop.API mcpManager *agentmcp.Manager mcpAPI *agentmcp.API @@ -474,6 +476,7 @@ func (a *agent) init() { return "" } a.processAPI = agentproc.NewAPI(a.logger.Named("processes"), a.execer, a.filesystem, pathStore, a.envInfo, a.updateCommandEnv, workingDirFn) + a.toolCalls = agenttoolcall.New(a.clock, a.processAPI.CancelToolCall) gitOpts := append([]agentgit.Option{agentgit.WithClock(a.clock)}, a.gitAPIOptions...) a.gitAPI = agentgit.NewAPI(a.logger.Named("git"), pathStore, gitOpts...) desktop := agentdesktop.NewPortableDesktop( diff --git a/agent/agent_test.go b/agent/agent_test.go index f15dc9f5f0..5248fc707c 100644 --- a/agent/agent_test.go +++ b/agent/agent_test.go @@ -4572,3 +4572,210 @@ func assertConnectionReport(t testing.TB, agentClient *agenttest.Client, report } assert.Equal(t, report.clientSessionID, reports[0].GetConnection().GetClientSessionId(), "connect reason should be %s", report.clientSessionID) } + +func TestAgent_ToolCall(t *testing.T) { + t.Parallel() + //nolint:dogsled + conn, _, _, agentFS, _ := setupAgent(t, agentsdk.Manifest{}, 0) + conn.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}}) + + // Every call uses the tool call context, as chatd does, so a method + // that wrongly sent the tool call ID would get a saved response. + tests := []struct { + name string + fs afero.Fs + act func(ctx context.Context, path string) (any, error) + saved func(workspacesdk.CancelToolCallResponse) (any, error) + }{ + { + name: "StartProcess", + fs: afero.NewOsFs(), + act: func(ctx context.Context, path string) (any, error) { + resp, err := conn.StartProcess(ctx, workspacesdk.StartProcessRequest{ + Command: "printf run > " + filepath.Base(path), + WorkDir: filepath.Dir(path), + }) + if err != nil { + return nil, err + } + out, err := conn.ProcessOutput(ctx, resp.ID, &workspacesdk.ProcessOutputOptions{Wait: true}) + if err != nil { + return nil, err + } + if out.ExitCode == nil { + return nil, xerrors.Errorf("process output reports no exit: %+v", out) + } + return resp, nil + }, + saved: func(r workspacesdk.CancelToolCallResponse) (any, error) { return r.StartProcessResult() }, + }, + { + name: "WriteFile", + fs: agentFS, + act: func(ctx context.Context, path string) (any, error) { + return nil, conn.WriteFile(ctx, path, strings.NewReader("run")) + }, + saved: func(r workspacesdk.CancelToolCallResponse) (any, error) { return nil, r.WriteFileResult() }, + }, + { + name: "EditFiles", + fs: agentFS, + act: func(ctx context.Context, path string) (any, error) { + return conn.EditFiles(ctx, workspacesdk.FileEditRequest{ + Files: []workspacesdk.FileEdits{{Path: path, Edits: []workspacesdk.FileEdit{{OldText: "reset", NewText: "run"}}}}, + IncludeDiff: true, + }) + }, + saved: func(r workspacesdk.CancelToolCallResponse) (any, error) { return r.EditFilesResult() }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + path := filepath.Join(t.TempDir(), "tool-call.txt") + reset := func() { require.NoError(t, afero.WriteFile(tc.fs, path, []byte("reset"), 0o600)) } + requireContent := func(want string) { + t.Helper() + content, err := afero.ReadFile(tc.fs, path) + require.NoError(t, err) + require.Equal(t, want, string(content)) + } + id := uuid.New() + toolCtx := workspacesdk.WithToolCallID(ctx, id) + + reset() + first, err := tc.act(toolCtx, path) + require.NoError(t, err) + requireContent("run") + reset() + repeat, err := tc.act(toolCtx, path) + require.NoError(t, err) + require.Equal(t, first, repeat) + requireContent("reset") + + canceled, err := conn.CancelToolCall(toolCtx, id) + require.NoError(t, err) + require.True(t, canceled.Received) + saved, err := tc.saved(canceled) + require.NoError(t, err) + require.Equal(t, first, saved) + + _, err = tc.act(toolCtx, path) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, http.StatusConflict, sdkErr.StatusCode()) + }) + } +} + +// TestAgent_ToolCallOtherChat checks that two chats using the same tool +// call ID get separate processes and cancels. +func TestAgent_ToolCallOtherChat(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + agentID := uuid.New() + //nolint:dogsled + connA, _, _, _, _ := setupAgent(t, agentsdk.Manifest{AgentID: agentID}, 0) + connA.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}}) + // connB shares connA's tailnet connection, which setupAgent closes. + connB := workspacesdk.NewAgentConn(connA.TailnetConn(), workspacesdk.AgentConnOptions{AgentID: agentID}) + connB.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}}) + + id := uuid.New() + toolCtx := workspacesdk.WithToolCallID(ctx, id) + _, err := connA.StartProcess(toolCtx, workspacesdk.StartProcessRequest{Command: "sleep 300"}) + require.NoError(t, err) + t.Cleanup(func() { + _, err := connA.CancelToolCall(workspacesdk.WithToolCallID(context.Background(), id), id) + assert.NoError(t, err) + }) + + _, err = connB.StartProcess(toolCtx, workspacesdk.StartProcessRequest{Command: "printf b"}) + require.NoError(t, err) + outB, err := connB.ProcessOutput(toolCtx, id.String(), &workspacesdk.ProcessOutputOptions{Wait: true}) + require.NoError(t, err) + require.Equal(t, "printf b", outB.Command) + require.Equal(t, "b", outB.Output) + + _, err = connB.CancelToolCall(toolCtx, id) + require.NoError(t, err) + // CancelToolCall marks the process canceled before returning, so a + // wrong cancel shows here without waiting. + outA, err := connA.ProcessOutput(toolCtx, id.String(), nil) + require.NoError(t, err) + require.Equal(t, "sleep 300", outA.Command) + require.False(t, outA.Canceled) +} + +func TestAgent_ProcessOutputTimeoutFromStartProcess(t *testing.T) { + t.Parallel() + //nolint:dogsled + conn, _, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0) + conn.SetExtraHeaders(http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}}) + + // A 1 ms timeout has passed before any wait uses it and one hour cannot + // pass during the test, so no case depends on timing. + tests := []struct { + name string + timeoutMs int64 + opts workspacesdk.ProcessOutputOptions + cancel bool + want workspacesdk.ProcessOutputResponse + }{ + { + name: "TimeoutFromStartProcess", + timeoutMs: 1, + opts: workspacesdk.ProcessOutputOptions{Wait: true, TimeoutFromStartProcess: true}, + want: workspacesdk.ProcessOutputResponse{Running: true, TimedOut: true}, + }, + { + name: "PlainWait", + timeoutMs: 1, + opts: workspacesdk.ProcessOutputOptions{Wait: true}, + cancel: true, + want: workspacesdk.ProcessOutputResponse{Canceled: true}, + }, + { + name: "BeforeTimeout", + timeoutMs: time.Hour.Milliseconds(), + want: workspacesdk.ProcessOutputResponse{Running: true}, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + id := uuid.New() + toolCtx := workspacesdk.WithToolCallID(ctx, id) + _, err := conn.StartProcess(toolCtx, workspacesdk.StartProcessRequest{Command: "sleep 300", TimeoutMs: tc.timeoutMs}) + require.NoError(t, err) + // Cancel the tool call so the agent does not wait for its process + // on close. + t.Cleanup(func() { + _, err := conn.CancelToolCall(workspacesdk.WithToolCallID(context.Background(), id), id) + assert.NoError(t, err) + }) + + type result struct { + resp workspacesdk.ProcessOutputResponse + err error + } + waited := make(chan result, 1) + go func() { + resp, err := conn.ProcessOutput(toolCtx, id.String(), &tc.opts) + waited <- result{resp, err} + }() + if tc.cancel { + _, err := conn.CancelToolCall(toolCtx, id) + require.NoError(t, err) + } + + got := testutil.RequireReceive(ctx, t, waited) + require.NoError(t, got.err) + require.Equal(t, tc.want.Running, got.resp.Running) + require.Equal(t, tc.want.TimedOut, got.resp.TimedOut) + require.Equal(t, tc.want.Canceled, got.resp.Canceled) + }) + } +} diff --git a/agent/agentchat/headers.go b/agent/agentchat/headers.go index 84db99bb25..90d42d7348 100644 --- a/agent/agentchat/headers.go +++ b/agent/agentchat/headers.go @@ -11,15 +11,20 @@ import ( // extractContext reads chat identity headers from the request. // Returns zero values if headers are absent (non-chat request). -func extractContext(r *http.Request) (chatID uuid.UUID, ancestorIDs []uuid.UUID, ok bool) { +func extractContext(r *http.Request) (chatCtx Context, ok bool) { raw := r.Header.Get(workspacesdk.CoderChatIDHeader) if raw == "" { - return uuid.Nil, nil, false + return Context{}, false } chatID, err := uuid.Parse(raw) if err != nil { - return uuid.Nil, nil, false + return Context{}, false } + toolCallID, err := uuid.Parse(r.Header.Get(workspacesdk.CoderToolCallIDHeader)) + if err != nil { + toolCallID = uuid.Nil + } + ancestorIDs := []uuid.UUID{} rawAncestors := r.Header.Get(workspacesdk.CoderAncestorChatIDsHeader) if rawAncestors != "" { var ids []string @@ -31,5 +36,5 @@ func extractContext(r *http.Request) (chatID uuid.UUID, ancestorIDs []uuid.UUID, } } } - return chatID, ancestorIDs, true + return Context{ID: chatID, AncestorIDs: ancestorIDs, ToolCallID: toolCallID}, true } diff --git a/agent/agentchat/log.go b/agent/agentchat/log.go index 319f6a79b6..764ca5b050 100644 --- a/agent/agentchat/log.go +++ b/agent/agentchat/log.go @@ -16,6 +16,7 @@ type chatContextKey struct{} type Context struct { ID uuid.UUID AncestorIDs []uuid.UUID + ToolCallID uuid.UUID // from Coder-Tool-Call-Id; uuid.Nil if absent or malformed } // FromContext returns the chat identity stored on the context. @@ -46,7 +47,7 @@ func Fields(ctx context.Context) []slog.Field { if !ok { return nil } - return chatFields(chatCtx.ID, chatCtx.AncestorIDs) + return chatFields(chatCtx) } // Middleware tags agent logs for requests that originate from @@ -55,30 +56,32 @@ func Fields(ctx context.Context) []slog.Field { // Install after loggermw.Logger so access-log enrichment can run. func Middleware(next http.Handler) http.Handler { return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { - chatID, ancestorIDs, ok := extractContext(r) + chatCtx, ok := extractContext(r) if !ok { next.ServeHTTP(rw, r) return } - fields := chatFields(chatID, ancestorIDs) + fields := chatFields(chatCtx) if requestLogger := loggermw.RequestLoggerFromContext(r.Context()); requestLogger != nil { requestLogger.WithFields(fields...) } - ctx := WithContext(r.Context(), chatID, ancestorIDs) - next.ServeHTTP(rw, r.WithContext(ctx)) + next.ServeHTTP(rw, r.WithContext(context.WithValue(r.Context(), chatContextKey{}, chatCtx))) }) } -func chatFields(chatID uuid.UUID, ancestorIDs []uuid.UUID) []slog.Field { - fields := []slog.Field{slog.F("chat_id", chatID.String())} - if len(ancestorIDs) == 0 { +func chatFields(chatCtx Context) []slog.Field { + fields := []slog.Field{slog.F("chat_id", chatCtx.ID.String())} + if chatCtx.ToolCallID != uuid.Nil { + fields = append(fields, slog.F("tool_call_id", chatCtx.ToolCallID.String())) + } + if len(chatCtx.AncestorIDs) == 0 { return fields } - ancestors := make([]string, 0, len(ancestorIDs)) - for _, id := range ancestorIDs { + ancestors := make([]string, 0, len(chatCtx.AncestorIDs)) + for _, id := range chatCtx.AncestorIDs { ancestors = append(ancestors, id.String()) } return append(fields, slog.F("ancestor_chat_ids", ancestors)) diff --git a/agent/agentchat/log_test.go b/agent/agentchat/log_test.go index 99cd94a133..ac2a53e784 100644 --- a/agent/agentchat/log_test.go +++ b/agent/agentchat/log_test.go @@ -21,6 +21,7 @@ func TestMiddlewareAccessLog(t *testing.T) { chatID := uuid.New() ancestorID := uuid.New() + toolCallID := uuid.New() sink := testutil.NewFakeSink(t) handler := tracing.StatusWriterMiddleware(loggermw.Logger(sink.Logger(), nil)( agentchat.Middleware(http.HandlerFunc(func(rw http.ResponseWriter, _ *http.Request) { @@ -31,6 +32,7 @@ func TestMiddlewareAccessLog(t *testing.T) { req := httptest.NewRequest(http.MethodGet, "/test", nil) req.Header.Set(workspacesdk.CoderChatIDHeader, chatID.String()) req.Header.Set(workspacesdk.CoderAncestorChatIDsHeader, mustMarshalJSON(t, []string{ancestorID.String()})) + req.Header.Set(workspacesdk.CoderToolCallIDHeader, toolCallID.String()) rw := httptest.NewRecorder() handler.ServeHTTP(rw, req) require.Equal(t, http.StatusNoContent, rw.Code) @@ -40,6 +42,7 @@ func TestMiddlewareAccessLog(t *testing.T) { fields := fieldsByName(entries[0].Fields) require.Equal(t, chatID.String(), fields["chat_id"]) require.Equal(t, []string{ancestorID.String()}, fields["ancestor_chat_ids"]) + require.Equal(t, toolCallID.String(), fields["tool_call_id"]) } func TestMiddlewareWithoutChatHeader(t *testing.T) { diff --git a/agent/agentproc/api.go b/agent/agentproc/api.go index b6f822e69e..5d9c0f6f99 100644 --- a/agent/agentproc/api.go +++ b/agent/agentproc/api.go @@ -7,6 +7,7 @@ import ( "fmt" "net/http" "sort" + "syscall" "time" "github.com/go-chi/chi/v5" @@ -81,13 +82,20 @@ func (api *API) handleStartProcess(rw http.ResponseWriter, r *http.Request) { }) return } - - var chatID string - if chatContext, ok := agentchat.FromContext(ctx); ok { - chatID = chatContext.ID.String() + if req.Background && req.TimeoutMs != 0 { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Background processes cannot have timeout_ms.", + }) + return } - proc, err := api.manager.start(req, chatID) + chat, _ := agentchat.FromContext(ctx) + id := chat.ToolCallID + if id == uuid.Nil { + id = uuid.New() + } + + proc, err := api.manager.start(req, chat.ID, id.String()) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to start process.", @@ -119,16 +127,12 @@ func (api *API) handleStartProcess(rw http.ResponseWriter, r *http.Request) { }) } -// handleListProcesses lists all tracked processes. +// handleListProcesses lists the processes of the requester's chat. func (api *API) handleListProcesses(rw http.ResponseWriter, r *http.Request) { ctx := r.Context() - var chatID string - if chatContext, ok := agentchat.FromContext(ctx); ok { - chatID = chatContext.ID.String() - } - - infos := api.manager.list(chatID) + chat, _ := agentchat.FromContext(ctx) + infos := api.manager.list(chat.ID) // Sort by running state (running first), then by started_at // descending so the most recent processes appear first. @@ -156,7 +160,8 @@ func (api *API) handleProcessOutput(rw http.ResponseWriter, r *http.Request) { logger := api.logger.With(agentchat.Fields(ctx)...) id := chi.URLParam(r, "id") - proc, ok := api.manager.get(id) + chat, _ := agentchat.FromContext(ctx) + proc, ok := api.manager.get(chat.ID, id) if !ok { httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{ Message: fmt.Sprintf("Process %q not found.", id), @@ -164,18 +169,6 @@ func (api *API) handleProcessOutput(rw http.ResponseWriter, r *http.Request) { return } - // Enforce chat ID isolation. If the request carries - // a chat context, only allow access to processes - // belonging to that chat. - if chatContext, ok := agentchat.FromContext(ctx); ok { - if proc.chatID != "" && proc.chatID != chatContext.ID.String() { - httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{ - Message: fmt.Sprintf("Process %q not found.", id), - }) - return - } - } - // Check for blocking mode via query params. waitStr := r.URL.Query().Get("wait") wantWait := waitStr == "true" @@ -195,8 +188,13 @@ func (api *API) handleProcessOutput(rw http.ResponseWriter, r *http.Request) { // Cap the wait at maxWaitDuration regardless of // client-supplied timeout. - waitCtx, waitCancel := context.WithTimeout(ctx, maxWaitDuration) + wait := maxWaitDuration + if r.URL.Query().Get("timeout_from_start_process") == "true" && !proc.waitUntil.IsZero() { + wait = min(wait, api.manager.clock.Until(proc.waitUntil)) + } + waitCtx, waitCancel := context.WithCancel(ctx) defer waitCancel() + defer api.manager.clock.AfterFunc(wait, waitCancel).Stop() _ = proc.waitForOutput(waitCtx) // Fall through to read snapshot below. @@ -216,26 +214,35 @@ func (api *API) handleProcessOutput(rw http.ResponseWriter, r *http.Request) { Running: info.Running, ExitCode: info.ExitCode, Command: info.Command, + TimedOut: info.Running && !proc.waitUntil.IsZero() && !api.manager.clock.Now().Before(proc.waitUntil), + Canceled: proc.canceled.Load(), }) } +// CancelToolCall cancels chat chatID's running process for tool call id: +// it marks the process canceled and sends it SIGKILL. +func (api *API) CancelToolCall(ctx context.Context, chatID, id uuid.UUID) { + proc, ok := api.manager.get(chatID, id.String()) + if !ok { + return + } + proc.mu.Lock() + defer proc.mu.Unlock() + if !proc.running { + return + } + proc.canceled.Store(true) + if err := signalProcess(proc.cmd.Process, syscall.SIGKILL); err != nil { + api.logger.Warn(ctx, "kill canceled tool call process", slog.F("process_id", id), slog.Error(err)) + } +} + // handleSignalProcess sends a signal to a running process. func (api *API) handleSignalProcess(rw http.ResponseWriter, r *http.Request) { ctx := r.Context() id := chi.URLParam(r, "id") - // Enforce chat ID isolation. - if chatContext, ok := agentchat.FromContext(ctx); ok { - proc, procOK := api.manager.get(id) - if procOK && proc.chatID != "" && proc.chatID != chatContext.ID.String() { - httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{ - Message: fmt.Sprintf("Process %q not found.", id), - }) - return - } - } - var req workspacesdk.SignalProcessRequest if err := json.NewDecoder(r.Body).Decode(&req); err != nil { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ @@ -262,7 +269,8 @@ func (api *API) handleSignalProcess(rw http.ResponseWriter, r *http.Request) { return } - if err := api.manager.signal(id, req.Signal); err != nil { + chat, _ := agentchat.FromContext(ctx) + if err := api.manager.signal(chat.ID, id, req.Signal); err != nil { switch { case errors.Is(err, errProcessNotFound): httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{ diff --git a/agent/agentproc/api_test.go b/agent/agentproc/api_test.go index 16b28f57d0..1f2c1878dc 100644 --- a/agent/agentproc/api_test.go +++ b/agent/agentproc/api_test.go @@ -45,6 +45,12 @@ func postStart(t *testing.T, handler http.Handler, req workspacesdk.StartProcess w := httptest.NewRecorder() r := httptest.NewRequestWithContext(ctx, http.MethodPost, "/start", bytes.NewReader(body)) + addHeaders(r, headers) + handler.ServeHTTP(w, r) + return w +} + +func addHeaders(r *http.Request, headers []http.Header) { for _, h := range headers { for k, vals := range h { for _, v := range vals { @@ -52,8 +58,6 @@ func postStart(t *testing.T, handler http.Handler, req workspacesdk.StartProcess } } } - handler.ServeHTTP(w, r) - return w } // getList sends a GET /list request and returns the recorder. @@ -101,7 +105,7 @@ func getOutputWithHeaders(t *testing.T, handler http.Handler, id string, headers // postSignal sends a POST /{id}/signal request and returns // the recorder. -func postSignal(t *testing.T, handler http.Handler, id string, req workspacesdk.SignalProcessRequest) *httptest.ResponseRecorder { +func postSignal(t *testing.T, handler http.Handler, id string, req workspacesdk.SignalProcessRequest, headers ...http.Header) *httptest.ResponseRecorder { t.Helper() ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) @@ -112,6 +116,7 @@ func postSignal(t *testing.T, handler http.Handler, id string, req workspacesdk. w := httptest.NewRecorder() r := httptest.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("/%s/signal", id), bytes.NewReader(body)) + addHeaders(r, headers) handler.ServeHTTP(w, r) return w } @@ -294,6 +299,43 @@ func TestStartProcess(t *testing.T) { require.Contains(t, resp.Message, "Command is required") }) + t.Run("BackgroundWithTimeout", func(t *testing.T) { + t.Parallel() + + handler := newTestAPI(t) + w := postStart(t, handler, workspacesdk.StartProcessRequest{ + Command: "echo background", + Background: true, + TimeoutMs: 1000, + }) + require.Equal(t, http.StatusBadRequest, w.Code) + + var resp codersdk.Response + require.NoError(t, json.NewDecoder(w.Body).Decode(&resp)) + require.Contains(t, resp.Message, "cannot have timeout_ms") + }) + + // Through the agent, the tool call middleware responds to a repeated + // start, so call the process API directly to test attaching. + t.Run("RepeatedToolCallAttaches", func(t *testing.T) { + t.Parallel() + + handler := newTestAPI(t) + id := uuid.NewString() + headers := http.Header{ + workspacesdk.CoderChatIDHeader: {uuid.NewString()}, + workspacesdk.CoderToolCallIDHeader: {id}, + } + req := workspacesdk.StartProcessRequest{Command: "sleep 300"} + require.Equal(t, http.StatusOK, postStart(t, handler, req, headers).Code) + + w := postStart(t, handler, req, headers) + require.Equal(t, http.StatusOK, w.Code) + var resp workspacesdk.StartProcessResponse + require.NoError(t, json.NewDecoder(w.Body).Decode(&resp)) + require.Equal(t, id, resp.ID) + }) + t.Run("MalformedJSON", func(t *testing.T) { t.Parallel() @@ -550,112 +592,6 @@ 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() @@ -868,34 +804,6 @@ func TestProcessOutput(t *testing.T) { require.Contains(t, resp.Message, "not found") }) - t.Run("ChatIDEnforcement", func(t *testing.T) { - t.Parallel() - - handler := newTestAPI(t) - - // Start a process with chat-a. - chatA := uuid.New() - id := startAndGetID(t, handler, workspacesdk.StartProcessRequest{ - Command: "echo secret", - Background: true, - }, http.Header{ - workspacesdk.CoderChatIDHeader: {chatA.String()}, - }) - waitForExit(t, handler, id) - - // Chat-b should NOT see this process. - chatB := uuid.New() - w1 := getOutputWithHeaders(t, handler, id, http.Header{ - workspacesdk.CoderChatIDHeader: {chatB.String()}, - }) - require.Equal(t, http.StatusNotFound, w1.Code) - - // Without any chat ID header, should return 200 - // (backwards compatible). - w2 := getOutput(t, handler, id) - require.Equal(t, http.StatusOK, w2.Code) - }) - t.Run("WaitForExit", func(t *testing.T) { t.Parallel() @@ -1006,17 +914,18 @@ func TestProcessOutput(t *testing.T) { }) } -func getOutputWithWait(t *testing.T, handler http.Handler, id string) *httptest.ResponseRecorder { +func getOutputWithWait(t *testing.T, handler http.Handler, id string, headers ...http.Header) *httptest.ResponseRecorder { t.Helper() ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong) defer cancel() - return getOutputWithWaitCtx(ctx, t, handler, id) + return getOutputWithWaitCtx(ctx, t, handler, id, headers...) } -func getOutputWithWaitCtx(ctx context.Context, t *testing.T, handler http.Handler, id string) *httptest.ResponseRecorder { +func getOutputWithWaitCtx(ctx context.Context, t *testing.T, handler http.Handler, id string, headers ...http.Header) *httptest.ResponseRecorder { t.Helper() path := fmt.Sprintf("/%s/output?wait=true", id) req := httptest.NewRequestWithContext(ctx, http.MethodGet, path, nil) + addHeaders(req, headers) w := httptest.NewRecorder() handler.ServeHTTP(w, req) return w @@ -1153,6 +1062,54 @@ func TestSignalProcess(t *testing.T) { }) } +func TestChatIsolation(t *testing.T) { + t.Parallel() + + chatA := http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}} + chatB := http.Header{workspacesdk.CoderChatIDHeader: {uuid.NewString()}} + noChat := http.Header{} + tests := []struct { + name string + owner http.Header + requester http.Header + listed bool + // The process has exited, so signaling a visible one returns 409. + wantOutput int + wantSignal int + }{ + {name: "SameChat", owner: chatA, requester: chatA, listed: true, wantOutput: http.StatusOK, wantSignal: http.StatusConflict}, + {name: "OtherChat", owner: chatA, requester: chatB, wantOutput: http.StatusNotFound, wantSignal: http.StatusNotFound}, + {name: "RequesterWithoutChat", owner: chatA, requester: noChat, wantOutput: http.StatusNotFound, wantSignal: http.StatusNotFound}, + {name: "OwnerWithoutChat", owner: noChat, requester: chatA, wantOutput: http.StatusNotFound, wantSignal: http.StatusNotFound}, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + handler := newTestAPI(t) + id := startAndGetID(t, handler, workspacesdk.StartProcessRequest{Command: "true"}, tc.owner) + require.Equal(t, http.StatusOK, getOutputWithWait(t, handler, id, tc.owner).Code) + + w := getListWithChatHeader(t, handler, tc.requester.Get(workspacesdk.CoderChatIDHeader)) + require.Equal(t, http.StatusOK, w.Code) + var list workspacesdk.ListProcessesResponse + require.NoError(t, json.NewDecoder(w.Body).Decode(&list)) + var wantIDs, gotIDs []string + if tc.listed { + wantIDs = []string{id} + } + for _, p := range list.Processes { + gotIDs = append(gotIDs, p.ID) + } + require.Equal(t, wantIDs, gotIDs) + + require.Equal(t, tc.wantOutput, getOutputWithHeaders(t, handler, id, tc.requester).Code) + w = postSignal(t, handler, id, workspacesdk.SignalProcessRequest{Signal: "kill"}, tc.requester) + require.Equal(t, tc.wantSignal, w.Code) + }) + } +} + func TestHandleStartProcess_ChatHeaders_EmptyWorkDir_StillNotifies(t *testing.T) { t.Parallel() diff --git a/agent/agentproc/process.go b/agent/agentproc/process.go index c5c93a2a1a..514c054c75 100644 --- a/agent/agentproc/process.go +++ b/agent/agentproc/process.go @@ -6,6 +6,7 @@ import ( "os" "os/exec" "sync" + "sync/atomic" "syscall" "time" @@ -25,8 +26,10 @@ var ( 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 + // kept before being automatically removed from the map, so + // that a chatd retry of its execute tool call can still read + // the output. + exitedProcessReapAge = time.Hour ) // process represents a running or completed process. @@ -36,12 +39,13 @@ type process struct { command string workDir string background bool - chatID string cmd *exec.Cmd cancel context.CancelFunc buf *HeadTailBuffer logger slog.Logger running bool + canceled atomic.Bool // set if its tool call is canceled while it runs + waitUntil time.Time // when TimeoutMs passes; zero if none exitCode *int startedAt int64 exitedAt *int64 @@ -71,6 +75,13 @@ func (p *process) output() (string, *workspacesdk.ProcessTruncation) { return p.buf.Output() } +// procKey identifies a process by chat ID and process ID. chatID is +// uuid.Nil for processes started without a chat. +type procKey struct { + chatID uuid.UUID + id string +} + // manager tracks processes spawned by the agent. type manager struct { mu sync.Mutex @@ -78,7 +89,7 @@ type manager struct { execer agentexec.Execer fs afero.Fs clock quartz.Clock - procs map[string]*process + procs map[procKey]*process closed bool updateEnv func(current []string) (updated []string, err error) workingDir func() string @@ -98,7 +109,7 @@ func newManager(logger slog.Logger, execer agentexec.Execer, fs afero.Fs, envInf execer: execer, fs: fs, clock: quartz.NewReal(), - procs: make(map[string]*process), + procs: make(map[procKey]*process), updateEnv: updateEnv, workingDir: workingDir, envInfo: envInfo, @@ -109,18 +120,27 @@ func newManager(logger slog.Logger, execer agentexec.Execer, fs afero.Fs, envInf // 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, chatID string) (*process, error) { +// +// A repeated start with the same chatID and id returns the existing +// process. The lookup and insert are not atomic; the tool call +// middleware serializes requests per tool call, and other starts get +// random IDs. +func (m *manager) start(req workspacesdk.StartProcessRequest, chatID uuid.UUID, id string) (*process, error) { + k := procKey{chatID: chatID, id: id} m.mu.Lock() if m.closed { m.mu.Unlock() return nil, xerrors.New("manager is closed") } + if proc, ok := m.procs[k]; ok { + m.mu.Unlock() + return proc, nil + } m.mu.Unlock() - id := uuid.New().String() logger := m.logger - if chatID != "" { - logger = logger.With(slog.F("chat_id", chatID)) + if chatID != uuid.Nil { + logger = logger.With(slog.F("chat_id", chatID.String())) } // Use a cancellable context so Close() can terminate @@ -167,7 +187,7 @@ func (m *manager) start(req workspacesdk.StartProcessRequest, chatID string) (*p } // Propagate the chat ID so child processes (e.g. // GIT_ASKPASS) can send it back to the server. - if chatID != "" { + if chatID != uuid.Nil { cmd.Env = append(cmd.Env, fmt.Sprintf("CODER_CHAT_ID=%s", chatID)) } @@ -176,21 +196,23 @@ func (m *manager) start(req workspacesdk.StartProcessRequest, chatID string) (*p return nil, xerrors.Errorf("start process: %w", err) } - now := m.clock.Now().Unix() + now := m.clock.Now() proc := &process{ id: id, command: req.Command, workDir: cmd.Dir, background: req.Background, - chatID: chatID, cmd: cmd, cancel: cancel, buf: buf, logger: logger, running: true, - startedAt: now, + startedAt: now.Unix(), done: make(chan struct{}), } + if req.TimeoutMs > 0 { + proc.waitUntil = now.Add(time.Duration(req.TimeoutMs) * time.Millisecond) + } m.mu.Lock() if m.closed { @@ -201,7 +223,7 @@ func (m *manager) start(req workspacesdk.StartProcessRequest, chatID string) (*p _ = cmd.Wait() return nil, xerrors.New("manager is closed") } - m.procs[id] = proc + m.procs[k] = proc m.mu.Unlock() go func() { @@ -240,37 +262,34 @@ func (m *manager) start(req workspacesdk.StartProcessRequest, chatID string) (*p return proc, nil } -// get returns a process by ID. -func (m *manager) get(id string) (*process, bool) { +// get returns chat chatID's process with ID id. +func (m *manager) get(chatID uuid.UUID, id string) (*process, bool) { m.mu.Lock() defer m.mu.Unlock() - proc, ok := m.procs[id] + proc, ok := m.procs[procKey{chatID: chatID, id: id}] return proc, ok } -// 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 { +// list returns info about chat chatID's processes. It also reaps +// processes of all chats that exited more than exitedProcessReapAge ago. +func (m *manager) list(chatID uuid.UUID) []workspacesdk.ProcessInfo { m.mu.Lock() defer m.mu.Unlock() now := m.clock.Now() infos := make([]workspacesdk.ProcessInfo, 0, len(m.procs)) - for id, proc := range m.procs { + for k, proc := range m.procs { info := proc.info() - // Reap processes that exited more than 5 minutes ago + // Reap processes that exited more than exitedProcessReapAge 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) + delete(m.procs, k) continue } } - // Filter by chatID if provided. - if chatID != "" && proc.chatID != chatID { + if k.chatID != chatID { continue } infos = append(infos, info) @@ -281,10 +300,8 @@ func (m *manager) list(chatID string) []workspacesdk.ProcessInfo { // signal sends a signal to a running process. It returns // sentinel errors errProcessNotFound and errProcessNotRunning // so callers can distinguish failure modes. -func (m *manager) signal(id string, sig string) error { - m.mu.Lock() - proc, ok := m.procs[id] - m.mu.Unlock() +func (m *manager) signal(chatID uuid.UUID, id string, sig string) error { + proc, ok := m.get(chatID, id) if !ok { return errProcessNotFound diff --git a/agent/agenttoolcall/agenttoolcall.go b/agent/agenttoolcall/agenttoolcall.go new file mode 100644 index 0000000000..b865b70b2f --- /dev/null +++ b/agent/agenttoolcall/agenttoolcall.go @@ -0,0 +1,183 @@ +// Package agenttoolcall deduplicates chat tool call requests and cancels +// tool calls. +package agenttoolcall + +import ( + "context" + "net/http" + "net/http/httptest" + "sync" + "time" + + "github.com/go-chi/chi/v5" + "github.com/google/uuid" + + "github.com/coder/coder/v2/agent/agentchat" + "github.com/coder/coder/v2/coderd/httpapi" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/workspacesdk" + "github.com/coder/quartz" +) + +// forgetAfter is how long an entry is kept after its last request or +// cancel. chatd resends unresolved tool calls on every retry, which keeps +// their entries. +const forgetAfter = time.Hour + +type key struct { + chatID uuid.UUID + id uuid.UUID +} + +type entry struct { + expiresAt time.Time + canceled bool + done chan struct{} // closed once resp is set; nil if a cancel created the entry + resp workspacesdk.CancelToolCallResponse +} + +// Table remembers tool calls by chat ID and tool call ID. +type Table struct { + clock quartz.Clock + cancel func(ctx context.Context, chatID, id uuid.UUID) + mu sync.Mutex + entries map[key]*entry +} + +// New returns a Table that calls cancel to stop work a canceled tool call +// left running, such as a process. +func New(clock quartz.Clock, cancel func(ctx context.Context, chatID, id uuid.UUID)) *Table { + return &Table{clock: clock, cancel: cancel, entries: make(map[key]*entry)} +} + +// add stores e under k and deletes expired entries. t.mu must be held. +func (t *Table) add(k key, e *entry) { + now := t.clock.Now() + e.expiresAt = now.Add(forgetAfter) + for k, old := range t.entries { + if now.After(old.expiresAt) { + delete(t.entries, k) + } + } + t.entries[k] = e +} + +// Middleware runs each tool call request once and responds to repeats +// with the saved response, or with 409 if the tool call was canceled. +func (t *Table) Middleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + chat, _ := agentchat.FromContext(ctx) + if chat.ToolCallID == uuid.Nil { + next.ServeHTTP(rw, r) + return + } + k := key{chatID: chat.ID, id: chat.ToolCallID} + + t.mu.Lock() + e, found := t.entries[k] + canceled := found && e.canceled + if found { + e.expiresAt = t.clock.Now().Add(forgetAfter) + } else { + e = &entry{done: make(chan struct{})} + t.add(k, e) + } + t.mu.Unlock() + + if canceled { + httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{ + Message: "Tool call was canceled.", + }) + return + } + if !found { + run(e, next, r) + } + select { + case <-e.done: + case <-ctx.Done(): + return + } + rw.Header().Set("Content-Type", e.resp.ContentType) + rw.WriteHeader(e.resp.Status) + _, _ = rw.Write(e.resp.Body) + }) +} + +// run runs next and saves its response in e. If next panics, run saves a +// 500 so that waiters are released. +func run(e *entry, next http.Handler, r *http.Request) { + rec := httptest.NewRecorder() + defer func() { + if e.resp.Status == 0 { + rec = httptest.NewRecorder() + httpapi.Write(r.Context(), rec, http.StatusInternalServerError, codersdk.Response{ + Message: "Tool call failed with an internal error. Its outcome is unknown.", + }) + e.resp = savedResponse(rec) + } + close(e.done) + }() + next.ServeHTTP(rec, r) + e.resp = savedResponse(rec) +} + +func savedResponse(rec *httptest.ResponseRecorder) workspacesdk.CancelToolCallResponse { + return workspacesdk.CancelToolCallResponse{ + Received: true, + Status: rec.Code, + ContentType: rec.Header().Get("Content-Type"), + Body: rec.Body.Bytes(), + } +} + +// Routes returns the handler for POST /{id}/cancel. +func (t *Table) Routes() http.Handler { + r := chi.NewRouter() + r.Post("/{id}/cancel", t.handleCancel) + return r +} + +// handleCancel marks the tool call canceled, waits for its request, stops +// its work, and responds with the saved response. It calls t.cancel even +// without an entry, because a process can outlive its entry. +func (t *Table) handleCancel(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + id, err := uuid.Parse(chi.URLParam(r, "id")) + if err != nil { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Invalid tool call ID.", + Detail: err.Error(), + }) + return + } + chat, _ := agentchat.FromContext(ctx) + k := key{chatID: chat.ID, id: id} + + t.mu.Lock() + e, found := t.entries[k] + if found { + e.canceled = true + e.expiresAt = t.clock.Now().Add(forgetAfter) + } else { + e = &entry{canceled: true} + t.add(k, e) + } + t.mu.Unlock() + + if e.done == nil { + t.cancel(ctx, chat.ID, id) + httpapi.Write(ctx, rw, http.StatusOK, workspacesdk.CancelToolCallResponse{}) + return + } + // Wait so that t.cancel also stops a process the request is still + // starting. + select { + case <-e.done: + case <-ctx.Done(): + return + } + t.cancel(ctx, chat.ID, id) + httpapi.Write(ctx, rw, http.StatusOK, e.resp) +} diff --git a/agent/agenttoolcall/agenttoolcall_test.go b/agent/agenttoolcall/agenttoolcall_test.go new file mode 100644 index 0000000000..eba788c1ea --- /dev/null +++ b/agent/agenttoolcall/agenttoolcall_test.go @@ -0,0 +1,272 @@ +package agenttoolcall_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strconv" + "sync/atomic" + "testing" + "time" + + "github.com/go-chi/chi/v5" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "cdr.dev/slog/v3/sloggers/slogtest" + "github.com/coder/coder/v2/agent/agentchat" + "github.com/coder/coder/v2/agent/agenttoolcall" + "github.com/coder/coder/v2/coderd/httpapi" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/codersdk/workspacesdk" + "github.com/coder/coder/v2/httpmw" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +// Per-route behavior is tested through the real agent in +// agent.TestAgent_ToolCall. + +type canceledCall struct { + chatID uuid.UUID + id uuid.UUID + runsFinished int64 // runs finished when cancel was called +} + +// testServer serves /run and /panic behind a Table and reports the Table's +// cancel calls on canceled. +type testServer struct { + clock *quartz.Mock + handler http.Handler + runs atomic.Int64 + finished atomic.Int64 + canceled chan canceledCall + entered chan struct{} // receives when a run starts + cancelEntered chan struct{} // receives when a cancel arrives, before the Table handles it + block chan struct{} // if set, runs wait until it is closed +} + +func newTestServer(t *testing.T) *testServer { + s := &testServer{ + clock: quartz.NewMock(t), + canceled: make(chan canceledCall, 10), + entered: make(chan struct{}, 10), + cancelEntered: make(chan struct{}, 10), + } + table := agenttoolcall.New(s.clock, func(_ context.Context, chatID, id uuid.UUID) { + s.canceled <- canceledCall{chatID: chatID, id: id, runsFinished: s.finished.Load()} + }) + r := chi.NewRouter() + r.Use(httpmw.Recover(slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})), agentchat.Middleware, table.Middleware) + r.Post("/run", func(rw http.ResponseWriter, r *http.Request) { + n := s.runs.Add(1) + s.entered <- struct{}{} + if s.block != nil { + <-s.block + } + httpapi.Write(r.Context(), rw, http.StatusCreated, codersdk.Response{Message: strconv.FormatInt(n, 10)}) + s.finished.Add(1) + }) + r.Post("/panic", func(http.ResponseWriter, *http.Request) { + s.runs.Add(1) + panic("handler panic") + }) + cancels := table.Routes() + r.Mount("/tool-calls", http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) { + s.cancelEntered <- struct{}{} + cancels.ServeHTTP(rw, r) + })) + s.handler = r + return s +} + +func (s *testServer) do(ctx context.Context, path string, chatID, id uuid.UUID) *httptest.ResponseRecorder { + req := httptest.NewRequestWithContext(ctx, http.MethodPost, path, nil) + req.Header.Set(workspacesdk.CoderChatIDHeader, chatID.String()) + if id != uuid.Nil { + req.Header.Set(workspacesdk.CoderToolCallIDHeader, id.String()) + } + rw := httptest.NewRecorder() + s.handler.ServeHTTP(rw, req) + return rw +} + +func (s *testServer) cancel(ctx context.Context, t *testing.T, chatID, id uuid.UUID) workspacesdk.CancelToolCallResponse { + t.Helper() + rw := s.do(ctx, "/tool-calls/"+id.String()+"/cancel", chatID, uuid.Nil) + require.Equal(t, http.StatusOK, rw.Code) + var resp workspacesdk.CancelToolCallResponse + require.NoError(t, json.NewDecoder(rw.Body).Decode(&resp)) + return resp +} + +func message(t *testing.T, body []byte) string { + t.Helper() + var resp codersdk.Response + require.NoError(t, json.Unmarshal(body, &resp)) + return resp.Message +} + +func TestMiddleware(t *testing.T) { + t.Parallel() + + chatA, chatB := uuid.New(), uuid.New() + tests := []struct { + name string + before func(ctx context.Context, t *testing.T, s *testServer, id uuid.UUID) + path string + wantCode int + wantMessage string + wantRuns int64 + }{ + { + name: "OtherChatsDoNotShare", + before: func(ctx context.Context, t *testing.T, s *testServer, id uuid.UUID) { + s.do(ctx, "/run", chatB, id) + s.cancel(ctx, t, chatB, id) + }, + path: "/run", + wantCode: http.StatusCreated, wantMessage: "2", wantRuns: 2, + }, + { + name: "ForgottenAfterOneHour", + before: func(ctx context.Context, _ *testing.T, s *testServer, id uuid.UUID) { + s.do(ctx, "/run", chatA, id) + s.clock.Advance(time.Hour + time.Second) + // Expired entries are deleted only when an entry is added. + s.do(ctx, "/run", chatA, uuid.New()) + }, + path: "/run", + wantCode: http.StatusCreated, wantMessage: "3", wantRuns: 3, + }, + { + name: "RepeatKeepsRecord", + before: func(ctx context.Context, _ *testing.T, s *testServer, id uuid.UUID) { + s.do(ctx, "/run", chatA, id) + s.clock.Advance(50 * time.Minute) + s.do(ctx, "/run", chatA, id) + s.clock.Advance(20 * time.Minute) + s.do(ctx, "/run", chatA, uuid.New()) + }, + path: "/run", + wantCode: http.StatusCreated, wantMessage: "1", wantRuns: 2, + }, + { + name: "CancelKeepsRecord", + before: func(ctx context.Context, t *testing.T, s *testServer, id uuid.UUID) { + s.cancel(ctx, t, chatA, id) + s.clock.Advance(50 * time.Minute) + s.cancel(ctx, t, chatA, id) + s.clock.Advance(20 * time.Minute) + s.do(ctx, "/run", chatA, uuid.New()) + }, + path: "/run", + wantCode: http.StatusConflict, wantMessage: "canceled", wantRuns: 1, + }, + { + name: "RefusedRequestKeepsRecord", + before: func(ctx context.Context, t *testing.T, s *testServer, id uuid.UUID) { + s.cancel(ctx, t, chatA, id) + s.clock.Advance(50 * time.Minute) + s.do(ctx, "/run", chatA, id) + s.clock.Advance(20 * time.Minute) + s.do(ctx, "/run", chatA, uuid.New()) + }, + path: "/run", + wantCode: http.StatusConflict, wantMessage: "canceled", wantRuns: 1, + }, + { + name: "RepeatAfterPanicGetsUnknownOutcome", + before: func(ctx context.Context, _ *testing.T, s *testServer, id uuid.UUID) { + s.do(ctx, "/panic", chatA, id) + }, + path: "/panic", + wantCode: http.StatusInternalServerError, wantMessage: "outcome is unknown", wantRuns: 1, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + s := newTestServer(t) + id := uuid.New() + tc.before(ctx, t, s, id) + rw := s.do(ctx, tc.path, chatA, id) + require.Equal(t, tc.wantCode, rw.Code) + require.Contains(t, message(t, rw.Body.Bytes()), tc.wantMessage) + require.Equal(t, tc.wantRuns, s.runs.Load()) + }) + } +} + +func TestConcurrentRequestsRunOnce(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + s := newTestServer(t) + s.block = make(chan struct{}) + chatID, id := uuid.New(), uuid.New() + + const requests = 10 + results := make(chan *httptest.ResponseRecorder, requests) + for range requests { + go func() { results <- s.do(ctx, "/run", chatID, id) }() + } + testutil.RequireReceive(ctx, t, s.entered) + close(s.block) + for range requests { + rw := testutil.RequireReceive(ctx, t, results) + require.Equal(t, http.StatusCreated, rw.Code) + require.Equal(t, "1", message(t, rw.Body.Bytes())) + } + require.EqualValues(t, 1, s.runs.Load()) +} + +func TestCancel(t *testing.T) { + t.Parallel() + + t.Run("NoRecord", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + s := newTestServer(t) + chatID, id := uuid.New(), uuid.New() + + require.False(t, s.cancel(ctx, t, chatID, id).Received) + c := testutil.RequireReceive(ctx, t, s.canceled) + require.Equal(t, chatID, c.chatID) + require.Equal(t, id, c.id) + require.Equal(t, http.StatusConflict, s.do(ctx, "/run", chatID, id).Code) + require.Zero(t, s.runs.Load()) + }) + + t.Run("DuringRun", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + s := newTestServer(t) + s.block = make(chan struct{}) + chatID, id := uuid.New(), uuid.New() + + runDone := make(chan *httptest.ResponseRecorder, 1) + go func() { runDone <- s.do(ctx, "/run", chatID, id) }() + testutil.RequireReceive(ctx, t, s.entered) + cancelDone := make(chan workspacesdk.CancelToolCallResponse, 1) + go func() { + rw := s.do(ctx, "/tool-calls/"+id.String()+"/cancel", chatID, uuid.Nil) + var resp workspacesdk.CancelToolCallResponse + _ = json.NewDecoder(rw.Body).Decode(&resp) + cancelDone <- resp + }() + // The cancel may mark the call before or after the run finishes; + // the assertions hold either way. + testutil.RequireReceive(ctx, t, s.cancelEntered) + close(s.block) + + run := testutil.RequireReceive(ctx, t, runDone) + resp := testutil.RequireReceive(ctx, t, cancelDone) + require.True(t, resp.Received) + require.Equal(t, run.Code, resp.Status) + require.Equal(t, run.Body.Bytes(), resp.Body) + c := testutil.RequireReceive(ctx, t, s.canceled) + require.EqualValues(t, 1, c.runsFinished, "the cancel hook runs after the run, so a process still starting is stopped") + }) +} diff --git a/agent/api.go b/agent/api.go index e290d34e85..8794c4995b 100644 --- a/agent/api.go +++ b/agent/api.go @@ -27,6 +27,7 @@ func (a *agent) apiHandler() http.Handler { tracing.Middleware(nil, []string{"/", "/api", "/api/**", "/debug/**"}, "agent"), loggermw.Logger(a.logger, nil), agentchat.Middleware, + a.toolCalls.Middleware, ) r.Get("/", func(rw http.ResponseWriter, r *http.Request) { httpapi.Write(r.Context(), rw, http.StatusOK, codersdk.Response{ @@ -37,6 +38,7 @@ func (a *agent) apiHandler() http.Handler { r.Mount("/api/v0", a.filesAPI.Routes()) r.Mount("/api/v0/git", a.gitAPI.Routes()) r.Mount("/api/v0/processes", a.processAPI.Routes()) + r.Mount("/api/v0/tool-calls", a.toolCalls.Routes()) r.Mount("/api/v0/desktop", a.desktopAPI.Routes()) r.Mount("/api/v0/mcp", a.mcpAPI.Routes()) r.Mount("/api/v0/context-config", a.contextConfigAPI.Routes()) diff --git a/codersdk/workspacesdk/agentconn.go b/codersdk/workspacesdk/agentconn.go index 8fdf80b790..0649e322b0 100644 --- a/codersdk/workspacesdk/agentconn.go +++ b/codersdk/workspacesdk/agentconn.go @@ -116,6 +116,7 @@ type AgentConn interface { RecreateDevcontainer(ctx context.Context, devcontainerID string) (codersdk.Response, error) SignalProcess(ctx context.Context, id string, signal string) error StartProcess(ctx context.Context, req StartProcessRequest) (StartProcessResponse, error) + CancelToolCall(ctx context.Context, id uuid.UUID) (CancelToolCallResponse, error) LS(ctx context.Context, path string, req LSRequest) (LSResponse, error) ResolvePath(ctx context.Context, path string) (string, error) ReadFile(ctx context.Context, path string, offset, limit int64) (io.ReadCloser, string, error) @@ -912,6 +913,10 @@ type StartProcessRequest struct { WorkDir string `json:"workdir,omitempty"` Env map[string]string `json:"env,omitempty"` Background bool `json:"background,omitempty"` + // TimeoutMs is the wait timeout in milliseconds after the process + // starts; see ProcessOutputOptions.TimeoutFromStartProcess. Not allowed + // with Background. + TimeoutMs int64 `json:"timeout_ms,omitempty"` } // StartProcessResponse is returned when a process is started. @@ -945,6 +950,8 @@ type ProcessOutputResponse struct { Running bool `json:"running"` ExitCode *int `json:"exit_code,omitempty"` Command string `json:"command,omitempty"` + TimedOut bool `json:"timed_out,omitempty"` // running past StartProcessRequest.TimeoutMs + Canceled bool `json:"canceled,omitempty"` // its tool call was canceled while it ran } // ProcessOutputOptions configures blocking behavior for @@ -953,6 +960,9 @@ type ProcessOutputOptions struct { // Wait enables blocking mode. When true, the request // blocks until the process exits or the context expires. Wait bool + // TimeoutFromStartProcess also ends a blocking wait at + // StartProcessRequest.TimeoutMs. + TimeoutFromStartProcess bool } // ProcessTruncation describes how process output was truncated. @@ -1117,13 +1127,17 @@ func (c *agentConn) WriteFile(ctx context.Context, path string, reader io.Reader ctx, span := tracing.StartSpan(ctx) defer span.End() - res, err := c.apiRequest(ctx, http.MethodPost, agentAPIPath("/api/v0/write-file", neturl.Values{ + res, err := c.toolCallRequest(ctx, http.MethodPost, agentAPIPath("/api/v0/write-file", neturl.Values{ "path": []string{path}, }), reader) if err != nil { return xerrors.Errorf("do request: %w", err) } defer res.Body.Close() + return readWriteFileResponse(res) +} + +func readWriteFileResponse(res *http.Response) error { if res.StatusCode != http.StatusOK { return codersdk.ReadBodyAsError(res) } @@ -1370,11 +1384,15 @@ type MCPToolContent struct { func (c *agentConn) StartProcess(ctx context.Context, req StartProcessRequest) (StartProcessResponse, error) { ctx, span := tracing.StartSpan(ctx) defer span.End() - res, err := c.apiRequest(ctx, http.MethodPost, "/api/v0/processes/start", req) + res, err := c.toolCallRequest(ctx, http.MethodPost, "/api/v0/processes/start", req) if err != nil { return StartProcessResponse{}, xerrors.Errorf("do request: %w", err) } defer res.Body.Close() + return readStartProcessResponse(res) +} + +func readStartProcessResponse(res *http.Response) (StartProcessResponse, error) { if res.StatusCode != http.StatusOK { return StartProcessResponse{}, codersdk.ReadBodyAsError(res) } @@ -1382,6 +1400,24 @@ func (c *agentConn) StartProcess(ctx context.Context, req StartProcessRequest) ( return resp, decodeAgentJSON(res, &resp) } +// CancelToolCall cancels tool call id of the connection's chat. The agent +// refuses later requests for it, waits for a running request, stops its +// process, and returns the saved response. +func (c *agentConn) CancelToolCall(ctx context.Context, id uuid.UUID) (CancelToolCallResponse, error) { + ctx, span := tracing.StartSpan(ctx) + defer span.End() + res, err := c.apiRequest(ctx, http.MethodPost, "/api/v0/tool-calls/"+id.String()+"/cancel", nil) + if err != nil { + return CancelToolCallResponse{}, xerrors.Errorf("do request: %w", err) + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + return CancelToolCallResponse{}, codersdk.ReadBodyAsError(res) + } + var resp CancelToolCallResponse + return resp, decodeAgentJSON(res, &resp) +} + // ListProcesses returns information about tracked processes on the agent. func (c *agentConn) ListProcesses(ctx context.Context) (ListProcessesResponse, error) { ctx, span := tracing.StartSpan(ctx) @@ -1436,11 +1472,14 @@ func (c *agentConn) CallMCPTool(ctx context.Context, req CallMCPToolRequest) (Ca func (c *agentConn) ProcessOutput(ctx context.Context, id string, opts *ProcessOutputOptions) (ProcessOutputResponse, error) { ctx, span := tracing.StartSpan(ctx) defer span.End() - path := "/api/v0/processes/" + id + "/output" + query := neturl.Values{} if opts != nil && opts.Wait { - path += "?wait=true" + query.Set("wait", "true") + if opts.TimeoutFromStartProcess { + query.Set("timeout_from_start_process", "true") + } } - res, err := c.apiRequest(ctx, http.MethodGet, path, nil) + res, err := c.apiRequest(ctx, http.MethodGet, agentAPIPath("/api/v0/processes/"+id+"/output", query), nil) if err != nil { return ProcessOutputResponse{}, xerrors.Errorf("do request: %w", err) } @@ -1478,11 +1517,15 @@ func (c *agentConn) EditFiles(ctx context.Context, edits FileEditRequest) (FileE ctx, span := tracing.StartSpan(ctx) defer span.End() - res, err := c.apiRequest(ctx, http.MethodPost, "/api/v0/edit-files", edits) + res, err := c.toolCallRequest(ctx, http.MethodPost, "/api/v0/edit-files", edits) if err != nil { return FileEditResponse{}, xerrors.Errorf("do request: %w", err) } defer res.Body.Close() + return readEditFilesResponse(res) +} + +func readEditFilesResponse(res *http.Response) (FileEditResponse, error) { if res.StatusCode != http.StatusOK { return FileEditResponse{}, codersdk.ReadBodyAsError(res) } @@ -1504,6 +1547,21 @@ func agentAPIPath(path string, query neturl.Values) string { // apiRequest makes a request to the workspace agent's HTTP API server. func (c *agentConn) apiRequest(ctx context.Context, method, path string, body interface{}) (*http.Response, error) { + return c.apiRequestWithHeader(ctx, method, path, body, nil) +} + +// toolCallRequest is apiRequest with CoderToolCallIDHeader from ctx. Use +// it only for requests that perform the tool call: the agent responds to +// any other request carrying the ID with the saved response. +func (c *agentConn) toolCallRequest(ctx context.Context, method, path string, body interface{}) (*http.Response, error) { + header := http.Header{} + if id, ok := ToolCallIDFromContext(ctx); ok { + header.Set(CoderToolCallIDHeader, id.String()) + } + return c.apiRequestWithHeader(ctx, method, path, body, header) +} + +func (c *agentConn) apiRequestWithHeader(ctx context.Context, method, path string, body interface{}, header http.Header) (*http.Response, error) { ctx, span := tracing.StartSpan(ctx) defer span.End() @@ -1543,6 +1601,9 @@ func (c *agentConn) apiRequest(ctx context.Context, method, path string, body in req.Header.Add(key, value) } } + for key, values := range header { + req.Header[key] = values + } return c.apiClient(ctx).Do(req) } diff --git a/codersdk/workspacesdk/agentconnmock/agentconnmock.go b/codersdk/workspacesdk/agentconnmock/agentconnmock.go index bbcc8fc526..86b52d4887 100644 --- a/codersdk/workspacesdk/agentconnmock/agentconnmock.go +++ b/codersdk/workspacesdk/agentconnmock/agentconnmock.go @@ -114,6 +114,21 @@ func (mr *MockAgentConnMockRecorder) CallMCPTool(ctx, req any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CallMCPTool", reflect.TypeOf((*MockAgentConn)(nil).CallMCPTool), ctx, req) } +// CancelToolCall mocks base method. +func (m *MockAgentConn) CancelToolCall(ctx context.Context, id uuid.UUID) (workspacesdk.CancelToolCallResponse, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "CancelToolCall", ctx, id) + ret0, _ := ret[0].(workspacesdk.CancelToolCallResponse) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// CancelToolCall indicates an expected call of CancelToolCall. +func (mr *MockAgentConnMockRecorder) CancelToolCall(ctx, id any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CancelToolCall", reflect.TypeOf((*MockAgentConn)(nil).CancelToolCall), ctx, id) +} + // Close mocks base method. func (m *MockAgentConn) Close() error { m.ctrl.T.Helper() diff --git a/codersdk/workspacesdk/toolcall.go b/codersdk/workspacesdk/toolcall.go new file mode 100644 index 0000000000..d81bf84f57 --- /dev/null +++ b/codersdk/workspacesdk/toolcall.go @@ -0,0 +1,64 @@ +package workspacesdk + +import ( + "bytes" + "context" + "io" + "net/http" + + "github.com/google/uuid" +) + +// CoderToolCallIDHeader carries the chat tool call ID. The agent +// deduplicates requests by chat and tool call ID, and uses the ID as the +// process ID of a process it starts. +const CoderToolCallIDHeader = "Coder-Tool-Call-Id" + +type toolCallIDContextKey struct{} + +// WithToolCallID returns a context whose StartProcess, EditFiles, and +// WriteFile requests send id in CoderToolCallIDHeader. The ID goes on the +// context because parallel tool calls share one AgentConn. +func WithToolCallID(ctx context.Context, id uuid.UUID) context.Context { + return context.WithValue(ctx, toolCallIDContextKey{}, id) +} + +// ToolCallIDFromContext returns the tool call ID set by WithToolCallID. +func ToolCallIDFromContext(ctx context.Context) (uuid.UUID, bool) { + id, ok := ctx.Value(toolCallIDContextKey{}).(uuid.UUID) + return id, ok +} + +// CancelToolCallResponse is the response to a tool call cancel. The agent +// has no record of a tool call it never received, or after it restarted +// or expired the record. +type CancelToolCallResponse struct { + Received bool `json:"received"` // false if the agent has no record of the tool call + Status int `json:"status,omitempty"` // saved response status, if Received + ContentType string `json:"content_type,omitempty"` // saved response Content-Type, if Received + Body []byte `json:"body,omitempty"` // saved response body, if Received +} + +// StartProcessResult decodes a saved start response as StartProcess +// does. +func (r CancelToolCallResponse) StartProcessResult() (StartProcessResponse, error) { + return readStartProcessResponse(r.httpResponse()) //nolint:bodyclose // In-memory body. +} + +// EditFilesResult decodes a saved edit response as EditFiles does. +func (r CancelToolCallResponse) EditFilesResult() (FileEditResponse, error) { + return readEditFilesResponse(r.httpResponse()) //nolint:bodyclose // In-memory body. +} + +// WriteFileResult decodes a saved write response as WriteFile does. +func (r CancelToolCallResponse) WriteFileResult() error { + return readWriteFileResponse(r.httpResponse()) //nolint:bodyclose // In-memory body. +} + +func (r CancelToolCallResponse) httpResponse() *http.Response { + return &http.Response{ + StatusCode: r.Status, + Header: http.Header{"Content-Type": {r.ContentType}}, + Body: io.NopCloser(bytes.NewReader(r.Body)), + } +}