mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: agents desktop recordings backend (#23894)
This PR introduces screen recording of the computer use agent using the virtual desktop. - Screen recording is triggered by a `wait_agent` tool call. Recording is stopped by a successful `wait_agent` tool call or when there hasn't been any desktop activity for 10 minutes. - Recordings are handled by the `portabledesktop` cli via the `record` command. The videos are sped up in periods of inactivity. - Recordings are saved to the database to the `chat_files` table. There's a hard limit of 100MB per recording. Larger recordings are dropped. - A successful `wait_agent` on a computer use subagent tool call returns a `recording_file_id`, later allowing the frontend to display the corresponding video.
This commit is contained in:
+1
-1
@@ -398,7 +398,7 @@ func (a *agent) init() {
|
||||
gitOpts := append([]agentgit.Option{agentgit.WithClock(a.clock)}, a.gitAPIOptions...)
|
||||
a.gitAPI = agentgit.NewAPI(a.logger.Named("git"), pathStore, gitOpts...)
|
||||
desktop := agentdesktop.NewPortableDesktop(
|
||||
a.logger.Named("desktop"), a.execer, a.scriptRunner.ScriptBinDir(),
|
||||
a.logger.Named("desktop"), a.execer, a.scriptRunner.ScriptBinDir(), nil,
|
||||
)
|
||||
a.desktopAPI = agentdesktop.NewAPI(a.logger.Named("desktop"), desktop, a.clock)
|
||||
a.mcpManager = agentmcp.NewManager(a.logger.Named("mcp"))
|
||||
|
||||
@@ -1,12 +1,17 @@
|
||||
package agentdesktop
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-chi/chi/v5"
|
||||
"github.com/google/uuid"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/agent/agentssh"
|
||||
@@ -47,6 +52,9 @@ type API struct {
|
||||
logger slog.Logger
|
||||
desktop Desktop
|
||||
clock quartz.Clock
|
||||
|
||||
closeMu sync.Mutex
|
||||
closed bool
|
||||
}
|
||||
|
||||
// NewAPI creates a new desktop streaming API.
|
||||
@@ -66,6 +74,10 @@ func (a *API) Routes() http.Handler {
|
||||
r := chi.NewRouter()
|
||||
r.Get("/vnc", a.handleDesktopVNC)
|
||||
r.Post("/action", a.handleAction)
|
||||
r.Route("/recording", func(r chi.Router) {
|
||||
r.Post("/start", a.handleRecordingStart)
|
||||
r.Post("/stop", a.handleRecordingStop)
|
||||
})
|
||||
return r
|
||||
}
|
||||
|
||||
@@ -116,6 +128,9 @@ func (a *API) handleAction(rw http.ResponseWriter, r *http.Request) {
|
||||
ctx := r.Context()
|
||||
handlerStart := a.clock.Now()
|
||||
|
||||
// Update last desktop action timestamp for idle recording monitor.
|
||||
a.desktop.RecordActivity()
|
||||
|
||||
// Ensure the desktop is running and grab native dimensions.
|
||||
cfg, err := a.desktop.Start(ctx)
|
||||
if err != nil {
|
||||
@@ -480,9 +495,150 @@ func (a *API) handleAction(rw http.ResponseWriter, r *http.Request) {
|
||||
|
||||
// Close shuts down the desktop session if one is running.
|
||||
func (a *API) Close() error {
|
||||
a.closeMu.Lock()
|
||||
if a.closed {
|
||||
a.closeMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
a.closed = true
|
||||
a.closeMu.Unlock()
|
||||
|
||||
return a.desktop.Close()
|
||||
}
|
||||
|
||||
// decodeRecordingRequest decodes and validates a recording request
|
||||
// from the HTTP body, returning the recording ID. Returns false if
|
||||
// the request was invalid and an error response was already written.
|
||||
func (*API) decodeRecordingRequest(rw http.ResponseWriter, r *http.Request) (string, bool) {
|
||||
ctx := r.Context()
|
||||
var req struct {
|
||||
RecordingID string `json:"recording_id"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: "Failed to decode request body.",
|
||||
Detail: err.Error(),
|
||||
})
|
||||
return "", false
|
||||
}
|
||||
if req.RecordingID == "" {
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: "Missing recording_id.",
|
||||
})
|
||||
return "", false
|
||||
}
|
||||
if _, err := uuid.Parse(req.RecordingID); err != nil {
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: "Invalid recording_id format.",
|
||||
Detail: "recording_id must be a valid UUID.",
|
||||
})
|
||||
return "", false
|
||||
}
|
||||
return req.RecordingID, true
|
||||
}
|
||||
|
||||
func (a *API) handleRecordingStart(rw http.ResponseWriter, r *http.Request) {
|
||||
ctx := r.Context()
|
||||
recordingID, ok := a.decodeRecordingRequest(rw, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
a.closeMu.Lock()
|
||||
if a.closed {
|
||||
a.closeMu.Unlock()
|
||||
httpapi.Write(ctx, rw, http.StatusServiceUnavailable, codersdk.Response{
|
||||
Message: "Desktop API is shutting down.",
|
||||
})
|
||||
return
|
||||
}
|
||||
a.closeMu.Unlock()
|
||||
|
||||
if err := a.desktop.StartRecording(ctx, recordingID); err != nil {
|
||||
if errors.Is(err, ErrDesktopClosed) {
|
||||
httpapi.Write(ctx, rw, http.StatusServiceUnavailable, codersdk.Response{
|
||||
Message: "Desktop API is shutting down.",
|
||||
})
|
||||
return
|
||||
}
|
||||
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
|
||||
Message: "Failed to start recording.",
|
||||
Detail: err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
httpapi.Write(ctx, rw, http.StatusOK, codersdk.Response{
|
||||
Message: "Recording started.",
|
||||
})
|
||||
}
|
||||
|
||||
func (a *API) handleRecordingStop(rw http.ResponseWriter, r *http.Request) {
|
||||
ctx := r.Context()
|
||||
|
||||
recordingID, ok := a.decodeRecordingRequest(rw, r)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
|
||||
a.closeMu.Lock()
|
||||
if a.closed {
|
||||
a.closeMu.Unlock()
|
||||
httpapi.Write(ctx, rw, http.StatusServiceUnavailable, codersdk.Response{
|
||||
Message: "Desktop API is shutting down.",
|
||||
})
|
||||
return
|
||||
}
|
||||
a.closeMu.Unlock()
|
||||
|
||||
// Stop recording (idempotent).
|
||||
// Use a context detached from the HTTP request so that if the
|
||||
// connection drops, the recording process can still shut down
|
||||
// gracefully. WithoutCancel preserves request-scoped values.
|
||||
stopCtx, stopCancel := context.WithTimeout(context.WithoutCancel(r.Context()), 30*time.Second)
|
||||
defer stopCancel()
|
||||
artifact, err := a.desktop.StopRecording(stopCtx, recordingID)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrUnknownRecording) {
|
||||
httpapi.Write(ctx, rw, http.StatusNotFound, codersdk.Response{
|
||||
Message: "Recording not found.",
|
||||
Detail: err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
if errors.Is(err, ErrRecordingCorrupted) {
|
||||
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
|
||||
Message: "Recording is corrupted.",
|
||||
Detail: err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
|
||||
Message: "Failed to stop recording.",
|
||||
Detail: err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
defer artifact.Reader.Close()
|
||||
|
||||
if artifact.Size > workspacesdk.MaxRecordingSize {
|
||||
a.logger.Warn(ctx, "recording file exceeds maximum size",
|
||||
slog.F("recording_id", recordingID),
|
||||
slog.F("size", artifact.Size),
|
||||
slog.F("max_size", workspacesdk.MaxRecordingSize),
|
||||
)
|
||||
httpapi.Write(ctx, rw, http.StatusRequestEntityTooLarge, codersdk.Response{
|
||||
Message: "Recording file exceeds maximum allowed size.",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
rw.Header().Set("Content-Type", "video/mp4")
|
||||
rw.Header().Set("Content-Length", strconv.FormatInt(artifact.Size, 10))
|
||||
rw.WriteHeader(http.StatusOK)
|
||||
_, _ = io.Copy(rw, artifact.Reader)
|
||||
}
|
||||
|
||||
// coordFromAction extracts the coordinate pair from a DesktopAction,
|
||||
// returning an error if the coordinate field is missing.
|
||||
func coordFromAction(action DesktopAction) (x, y int, err error) {
|
||||
|
||||
@@ -4,12 +4,17 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"slices"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/xerrors"
|
||||
@@ -21,6 +26,16 @@ import (
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
// Test recording UUIDs used across tests.
|
||||
const (
|
||||
testRecIDDefault = "870e1f02-8118-4300-a37e-4adb0117baf3"
|
||||
testRecIDStartIdempotent = "250a2ffb-a5e5-4c94-9754-4d6a4ab7ba20"
|
||||
testRecIDStopIdempotent = "38f8a378-f98f-4758-a4ae-950b44cf989a"
|
||||
testRecIDConcurrentA = "8dc173eb-23c6-4601-a485-b6dfb2a42c3a"
|
||||
testRecIDConcurrentB = "fea490d4-70f0-4798-a181-29d65ce25ae1"
|
||||
testRecIDRestart = "75173a0d-b018-4e2e-a771-defa3fc6af69"
|
||||
)
|
||||
|
||||
// Ensure fakeDesktop satisfies the Desktop interface at compile time.
|
||||
var _ agentdesktop.Desktop = (*fakeDesktop)(nil)
|
||||
|
||||
@@ -43,6 +58,14 @@ type fakeDesktop struct {
|
||||
lastTyped string
|
||||
lastKeyDown string
|
||||
lastKeyUp string
|
||||
|
||||
// Recording tracking (guarded by recMu).
|
||||
recMu sync.Mutex
|
||||
recordings map[string]string // ID → file path
|
||||
stopCalls []string // recording IDs passed to StopRecording
|
||||
recStopCh chan string // optional: signaled when StopRecording is called
|
||||
startCount int // incremented on each new recording start
|
||||
activityCount int // incremented by RecordActivity
|
||||
}
|
||||
|
||||
func (f *fakeDesktop) Start(context.Context) (agentdesktop.DisplayConfig, error) {
|
||||
@@ -107,11 +130,140 @@ func (f *fakeDesktop) CursorPosition(context.Context) (x int, y int, err error)
|
||||
return f.cursorPos[0], f.cursorPos[1], nil
|
||||
}
|
||||
|
||||
func (f *fakeDesktop) StartRecording(_ context.Context, recordingID string) error {
|
||||
f.recMu.Lock()
|
||||
defer f.recMu.Unlock()
|
||||
if f.recordings == nil {
|
||||
f.recordings = make(map[string]string)
|
||||
}
|
||||
if path, ok := f.recordings[recordingID]; ok {
|
||||
// Check if already stopped (file still exists but stop was
|
||||
// called). For the fake, a stopped recording means its ID
|
||||
// appears in stopCalls. In that case, remove the old file
|
||||
// and start fresh.
|
||||
stopped := slices.Contains(f.stopCalls, recordingID)
|
||||
if !stopped {
|
||||
// Active recording - no-op.
|
||||
return nil
|
||||
}
|
||||
// Completed recording - discard old file, start fresh.
|
||||
_ = os.Remove(path)
|
||||
delete(f.recordings, recordingID)
|
||||
}
|
||||
f.startCount++
|
||||
tmpFile, err := os.CreateTemp("", "fake-recording-*.mp4")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, _ = tmpFile.Write([]byte(fmt.Sprintf("fake-mp4-data-%s-%d", recordingID, f.startCount)))
|
||||
_ = tmpFile.Close()
|
||||
f.recordings[recordingID] = tmpFile.Name()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeDesktop) StopRecording(_ context.Context, recordingID string) (*agentdesktop.RecordingArtifact, error) {
|
||||
f.recMu.Lock()
|
||||
defer f.recMu.Unlock()
|
||||
if f.recordings == nil {
|
||||
return nil, agentdesktop.ErrUnknownRecording
|
||||
}
|
||||
path, ok := f.recordings[recordingID]
|
||||
if !ok {
|
||||
return nil, agentdesktop.ErrUnknownRecording
|
||||
}
|
||||
f.stopCalls = append(f.stopCalls, recordingID)
|
||||
if f.recStopCh != nil {
|
||||
select {
|
||||
case f.recStopCh <- recordingID:
|
||||
default:
|
||||
}
|
||||
}
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
info, err := file.Stat()
|
||||
if err != nil {
|
||||
_ = file.Close()
|
||||
return nil, err
|
||||
}
|
||||
return &agentdesktop.RecordingArtifact{
|
||||
Reader: file,
|
||||
Size: info.Size(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (f *fakeDesktop) RecordActivity() {
|
||||
f.recMu.Lock()
|
||||
f.activityCount++
|
||||
f.recMu.Unlock()
|
||||
}
|
||||
|
||||
func (f *fakeDesktop) Close() error {
|
||||
f.closed = true
|
||||
f.recMu.Lock()
|
||||
defer f.recMu.Unlock()
|
||||
for _, path := range f.recordings {
|
||||
_ = os.Remove(path)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// failStartRecordingDesktop wraps fakeDesktop and overrides
|
||||
// StartRecording to always return an error.
|
||||
type failStartRecordingDesktop struct {
|
||||
fakeDesktop
|
||||
startRecordingErr error
|
||||
}
|
||||
|
||||
func (f *failStartRecordingDesktop) StartRecording(_ context.Context, _ string) error {
|
||||
return f.startRecordingErr
|
||||
}
|
||||
|
||||
// corruptedStopDesktop wraps fakeDesktop and overrides
|
||||
// StopRecording to always return ErrRecordingCorrupted.
|
||||
type corruptedStopDesktop struct {
|
||||
fakeDesktop
|
||||
}
|
||||
|
||||
func (*corruptedStopDesktop) StopRecording(_ context.Context, _ string) (*agentdesktop.RecordingArtifact, error) {
|
||||
return nil, agentdesktop.ErrRecordingCorrupted
|
||||
}
|
||||
|
||||
// oversizedFakeDesktop wraps fakeDesktop and expands recording files
|
||||
// beyond MaxRecordingSize when StopRecording is called.
|
||||
type oversizedFakeDesktop struct {
|
||||
fakeDesktop
|
||||
}
|
||||
|
||||
func (f *oversizedFakeDesktop) StopRecording(ctx context.Context, recordingID string) (*agentdesktop.RecordingArtifact, error) {
|
||||
artifact, err := f.fakeDesktop.StopRecording(ctx, recordingID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Close the original reader since we're going to re-open after truncation.
|
||||
artifact.Reader.Close()
|
||||
|
||||
// Look up the path from the fakeDesktop recordings.
|
||||
f.fakeDesktop.recMu.Lock()
|
||||
path := f.fakeDesktop.recordings[recordingID]
|
||||
f.fakeDesktop.recMu.Unlock()
|
||||
|
||||
// Expand the file to exceed the maximum recording size.
|
||||
if err := os.Truncate(path, workspacesdk.MaxRecordingSize+1); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Re-open the truncated file.
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &agentdesktop.RecordingArtifact{
|
||||
Reader: file,
|
||||
Size: workspacesdk.MaxRecordingSize + 1,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func TestHandleDesktopVNC_StartError(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -134,6 +286,37 @@ func TestHandleDesktopVNC_StartError(t *testing.T) {
|
||||
assert.Equal(t, "Failed to start desktop session.", resp.Message)
|
||||
}
|
||||
|
||||
func TestHandleAction_CallsRecordActivity(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
body := agentdesktop.DesktopAction{
|
||||
Action: "left_click",
|
||||
Coordinate: &[2]int{100, 200},
|
||||
}
|
||||
b, err := json.Marshal(body)
|
||||
require.NoError(t, err)
|
||||
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/action", bytes.NewReader(b))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
handler := api.Routes()
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
|
||||
fake.recMu.Lock()
|
||||
count := fake.activityCount
|
||||
fake.recMu.Unlock()
|
||||
assert.Equal(t, 1, count, "handleAction should call RecordActivity exactly once")
|
||||
}
|
||||
|
||||
func TestHandleAction_Screenshot(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -574,3 +757,481 @@ func TestHandleAction_CursorPositionReturnsDeclaredCoordinates(t *testing.T) {
|
||||
// Native (960,540) in 1920x1080 should map to declared space in 1280x720.
|
||||
assert.Equal(t, "x=640,y=360", resp.Output)
|
||||
}
|
||||
|
||||
func TestRecordingStartStop(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
// Start recording.
|
||||
startBody, err := json.Marshal(map[string]string{"recording_id": testRecIDDefault})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(startBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
|
||||
// Stop recording.
|
||||
stopBody, err := json.Marshal(map[string]string{"recording_id": testRecIDDefault})
|
||||
require.NoError(t, err)
|
||||
rr = httptest.NewRecorder()
|
||||
req = httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader(stopBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
assert.Equal(t, "video/mp4", rr.Header().Get("Content-Type"))
|
||||
assert.Equal(t, []byte("fake-mp4-data-"+testRecIDDefault+"-1"), rr.Body.Bytes())
|
||||
}
|
||||
|
||||
func TestRecordingStartFails(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &failStartRecordingDesktop{
|
||||
fakeDesktop: fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
},
|
||||
startRecordingErr: xerrors.New("start recording error"),
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
body, err := json.Marshal(map[string]string{"recording_id": uuid.New().String()})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(body))
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
assert.Equal(t, http.StatusInternalServerError, rr.Code)
|
||||
|
||||
var resp codersdk.Response
|
||||
err = json.NewDecoder(rr.Body).Decode(&resp)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Failed to start recording.", resp.Message)
|
||||
}
|
||||
|
||||
func TestRecordingStartIdempotent(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
// Start same recording twice - both should succeed.
|
||||
for range 2 {
|
||||
body, err := json.Marshal(map[string]string{"recording_id": testRecIDStartIdempotent})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(body))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
}
|
||||
|
||||
// Stop once, verify normal response.
|
||||
stopBody, err := json.Marshal(map[string]string{"recording_id": testRecIDStartIdempotent})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader(stopBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
assert.Equal(t, "video/mp4", rr.Header().Get("Content-Type"))
|
||||
assert.Equal(t, []byte("fake-mp4-data-"+testRecIDStartIdempotent+"-1"), rr.Body.Bytes())
|
||||
}
|
||||
|
||||
func TestRecordingStopIdempotent(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
// Start recording.
|
||||
startBody, err := json.Marshal(map[string]string{"recording_id": testRecIDStopIdempotent})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(startBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
|
||||
// Stop twice - both should succeed with identical data.
|
||||
var bodies [2][]byte
|
||||
for i := range 2 {
|
||||
body, err := json.Marshal(map[string]string{"recording_id": testRecIDStopIdempotent})
|
||||
require.NoError(t, err)
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader(body))
|
||||
handler.ServeHTTP(recorder, request)
|
||||
require.Equal(t, http.StatusOK, recorder.Code)
|
||||
assert.Equal(t, "video/mp4", recorder.Header().Get("Content-Type"))
|
||||
bodies[i] = recorder.Body.Bytes()
|
||||
}
|
||||
assert.Equal(t, bodies[0], bodies[1])
|
||||
}
|
||||
|
||||
func TestRecordingStopInvalidIDFormat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
body, err := json.Marshal(map[string]string{"recording_id": "not-a-uuid"})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader(body))
|
||||
handler.ServeHTTP(rr, req)
|
||||
assert.Equal(t, http.StatusBadRequest, rr.Code)
|
||||
}
|
||||
|
||||
func TestRecordingStopUnknownRecording(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
// Send a valid UUID that was never started - should reach
|
||||
// StopRecording, get ErrUnknownRecording, and return 404.
|
||||
body, err := json.Marshal(map[string]string{"recording_id": uuid.New().String()})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader(body))
|
||||
handler.ServeHTTP(rr, req)
|
||||
assert.Equal(t, http.StatusNotFound, rr.Code)
|
||||
|
||||
var resp codersdk.Response
|
||||
err = json.NewDecoder(rr.Body).Decode(&resp)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Recording not found.", resp.Message)
|
||||
}
|
||||
|
||||
func TestRecordingStopOversizedFile(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &oversizedFakeDesktop{
|
||||
fakeDesktop: fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
// Start recording.
|
||||
recID := uuid.New().String()
|
||||
startBody, err := json.Marshal(map[string]string{"recording_id": recID})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(startBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
|
||||
// Stop recording - file exceeds max size, expect 413.
|
||||
stopBody, err := json.Marshal(map[string]string{"recording_id": recID})
|
||||
require.NoError(t, err)
|
||||
rr = httptest.NewRecorder()
|
||||
req = httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader(stopBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
assert.Equal(t, http.StatusRequestEntityTooLarge, rr.Code)
|
||||
|
||||
var resp codersdk.Response
|
||||
err = json.NewDecoder(rr.Body).Decode(&resp)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Recording file exceeds maximum allowed size.", resp.Message)
|
||||
}
|
||||
|
||||
func TestRecordingMultipleSimultaneous(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
// Start two recordings with different IDs.
|
||||
for _, id := range []string{testRecIDConcurrentA, testRecIDConcurrentB} {
|
||||
body, err := json.Marshal(map[string]string{"recording_id": id})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(body))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
}
|
||||
|
||||
// Stop both and verify each returns its own data.
|
||||
expected := map[string][]byte{
|
||||
testRecIDConcurrentA: []byte("fake-mp4-data-" + testRecIDConcurrentA + "-1"),
|
||||
testRecIDConcurrentB: []byte("fake-mp4-data-" + testRecIDConcurrentB + "-2"),
|
||||
}
|
||||
for _, id := range []string{testRecIDConcurrentA, testRecIDConcurrentB} {
|
||||
body, err := json.Marshal(map[string]string{"recording_id": id})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader(body))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
assert.Equal(t, "video/mp4", rr.Header().Get("Content-Type"))
|
||||
assert.Equal(t, expected[id], rr.Body.Bytes())
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordingStartMalformedBody(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader([]byte("not json")))
|
||||
handler.ServeHTTP(rr, req)
|
||||
assert.Equal(t, http.StatusBadRequest, rr.Code)
|
||||
}
|
||||
|
||||
func TestRecordingStartEmptyID(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
body, err := json.Marshal(map[string]string{"recording_id": ""})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(body))
|
||||
handler.ServeHTTP(rr, req)
|
||||
assert.Equal(t, http.StatusBadRequest, rr.Code)
|
||||
}
|
||||
|
||||
func TestRecordingStopEmptyID(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
body, err := json.Marshal(map[string]string{"recording_id": ""})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader(body))
|
||||
handler.ServeHTTP(rr, req)
|
||||
assert.Equal(t, http.StatusBadRequest, rr.Code)
|
||||
}
|
||||
|
||||
func TestRecordingStopMalformedBody(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader([]byte("not json")))
|
||||
handler.ServeHTTP(rr, req)
|
||||
assert.Equal(t, http.StatusBadRequest, rr.Code)
|
||||
}
|
||||
|
||||
func TestRecordingStartAfterCompleted(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
// Step 1: Start recording.
|
||||
startBody, err := json.Marshal(map[string]string{"recording_id": testRecIDRestart})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(startBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
|
||||
// Step 2: Stop recording (gets first MP4 data).
|
||||
stopBody, err := json.Marshal(map[string]string{"recording_id": testRecIDRestart})
|
||||
require.NoError(t, err)
|
||||
rr = httptest.NewRecorder()
|
||||
req = httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader(stopBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
assert.Equal(t, "video/mp4", rr.Header().Get("Content-Type"))
|
||||
firstData := rr.Body.Bytes()
|
||||
require.NotEmpty(t, firstData)
|
||||
|
||||
// Step 3: Start again with the same ID - should succeed
|
||||
// (old file discarded, new recording started).
|
||||
rr = httptest.NewRecorder()
|
||||
req = httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(startBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
|
||||
// Step 4: Stop again - should return NEW MP4 data.
|
||||
rr = httptest.NewRecorder()
|
||||
req = httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader(stopBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
assert.Equal(t, "video/mp4", rr.Header().Get("Content-Type"))
|
||||
secondData := rr.Body.Bytes()
|
||||
require.NotEmpty(t, secondData)
|
||||
|
||||
// The two recordings should have different data because the
|
||||
// fake increments a counter on each fresh start.
|
||||
assert.NotEqual(t, firstData, secondData,
|
||||
"restarted recording should produce different data")
|
||||
}
|
||||
|
||||
func TestRecordingStartAfterClose(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
// Close the API before sending the request.
|
||||
api.Close()
|
||||
|
||||
body, err := json.Marshal(map[string]string{"recording_id": uuid.New().String()})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(body))
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
assert.Equal(t, http.StatusServiceUnavailable, rr.Code)
|
||||
|
||||
var resp codersdk.Response
|
||||
err = json.NewDecoder(rr.Body).Decode(&resp)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Desktop API is shutting down.", resp.Message)
|
||||
}
|
||||
|
||||
func TestRecordingStartDesktopClosed(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
// StartRecording returns ErrDesktopClosed to simulate a race
|
||||
// where the desktop is closed between the API-level check and
|
||||
// the desktop-level StartRecording call.
|
||||
fake := &failStartRecordingDesktop{
|
||||
fakeDesktop: fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
},
|
||||
startRecordingErr: agentdesktop.ErrDesktopClosed,
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
body, err := json.Marshal(map[string]string{"recording_id": uuid.New().String()})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(body))
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
assert.Equal(t, http.StatusServiceUnavailable, rr.Code)
|
||||
|
||||
var resp codersdk.Response
|
||||
err = json.NewDecoder(rr.Body).Decode(&resp)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Desktop API is shutting down.", resp.Message)
|
||||
}
|
||||
|
||||
func TestRecordingStopCorrupted(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
fake := &corruptedStopDesktop{
|
||||
fakeDesktop: fakeDesktop{
|
||||
startCfg: agentdesktop.DisplayConfig{Width: 1920, Height: 1080},
|
||||
},
|
||||
}
|
||||
api := agentdesktop.NewAPI(logger, fake, nil)
|
||||
defer api.Close()
|
||||
|
||||
handler := api.Routes()
|
||||
|
||||
// Start a recording so the stop has something to find.
|
||||
recID := uuid.New().String()
|
||||
startBody, err := json.Marshal(map[string]string{"recording_id": recID})
|
||||
require.NoError(t, err)
|
||||
rr := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/recording/start", bytes.NewReader(startBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
require.Equal(t, http.StatusOK, rr.Code)
|
||||
|
||||
// Stop returns ErrRecordingCorrupted.
|
||||
stopBody, err := json.Marshal(map[string]string{"recording_id": recID})
|
||||
require.NoError(t, err)
|
||||
rr = httptest.NewRecorder()
|
||||
req = httptest.NewRequest(http.MethodPost, "/recording/stop", bytes.NewReader(stopBody))
|
||||
handler.ServeHTTP(rr, req)
|
||||
|
||||
assert.Equal(t, http.StatusInternalServerError, rr.Code)
|
||||
|
||||
var respStop codersdk.Response
|
||||
err = json.NewDecoder(rr.Body).Decode(&respStop)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "Recording is corrupted.", respStop.Message)
|
||||
}
|
||||
|
||||
@@ -2,7 +2,10 @@ package agentdesktop
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net"
|
||||
|
||||
"golang.org/x/xerrors"
|
||||
)
|
||||
|
||||
// Desktop abstracts a virtual desktop session running inside a workspace.
|
||||
@@ -58,10 +61,52 @@ type Desktop interface {
|
||||
// CursorPosition returns the current cursor coordinates.
|
||||
CursorPosition(ctx context.Context) (x, y int, err error)
|
||||
|
||||
// RecordActivity marks the desktop as having received user
|
||||
// interaction, resetting the idle-recording timer.
|
||||
RecordActivity()
|
||||
|
||||
// StartRecording begins recording the desktop to an MP4 file
|
||||
// using the caller-provided recording ID. Safe to call
|
||||
// repeatedly - active recordings continue unchanged, stopped
|
||||
// recordings are discarded and restarted. Concurrent recordings
|
||||
// are supported.
|
||||
StartRecording(ctx context.Context, recordingID string) error
|
||||
|
||||
// StopRecording finalizes the recording identified by the given
|
||||
// ID. Idempotent - safe to call on an already-stopped recording.
|
||||
// Returns a RecordingArtifact that the caller can stream. The
|
||||
// caller must close the artifact when done. Returns an error if
|
||||
// the recording ID is unknown.
|
||||
StopRecording(ctx context.Context, recordingID string) (*RecordingArtifact, error)
|
||||
|
||||
// Close shuts down the desktop session and cleans up resources.
|
||||
Close() error
|
||||
}
|
||||
|
||||
// ErrUnknownRecording is returned by StopRecording when the
|
||||
// recording ID is not recognized.
|
||||
var ErrUnknownRecording = xerrors.New("unknown recording ID")
|
||||
|
||||
// ErrDesktopClosed is returned when an operation is attempted on a
|
||||
// closed desktop session.
|
||||
var ErrDesktopClosed = xerrors.New("desktop closed")
|
||||
|
||||
// ErrRecordingCorrupted is returned by StopRecording when the
|
||||
// recording process was force-killed and the artifact is likely
|
||||
// incomplete or corrupt.
|
||||
var ErrRecordingCorrupted = xerrors.New("recording corrupted: process was force-killed")
|
||||
|
||||
// RecordingArtifact is a finalized recording returned by StopRecording.
|
||||
// The caller streams the artifact and must call Close when done. The
|
||||
// artifact remains valid even if the same recording ID is restarted
|
||||
// or the desktop is closed while the caller is reading.
|
||||
type RecordingArtifact struct {
|
||||
// Reader is the MP4 content. Callers must close it when done.
|
||||
Reader io.ReadCloser
|
||||
// Size is the byte length of the MP4 content.
|
||||
Size int64
|
||||
}
|
||||
|
||||
// DisplayConfig describes a running desktop session.
|
||||
type DisplayConfig struct {
|
||||
Width int // native width in pixels
|
||||
|
||||
@@ -3,6 +3,7 @@ package agentdesktop
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
@@ -11,6 +12,7 @@ import (
|
||||
"runtime"
|
||||
"strconv"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"golang.org/x/xerrors"
|
||||
@@ -18,6 +20,7 @@ import (
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/agent/agentexec"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
// portableDesktopOutput is the JSON output from
|
||||
@@ -49,32 +52,65 @@ type screenshotOutput struct {
|
||||
Data string `json:"data"`
|
||||
}
|
||||
|
||||
// recordingProcess tracks a single desktop recording subprocess.
|
||||
type recordingProcess struct {
|
||||
cmd *exec.Cmd
|
||||
filePath string
|
||||
stopped bool
|
||||
killed bool // true when the process was SIGKILLed
|
||||
done chan struct{} // closed when cmd.Wait() returns
|
||||
waitErr error // set before done is closed
|
||||
stopOnce sync.Once
|
||||
idleCancel context.CancelFunc // cancels the per-recording idle goroutine
|
||||
idleDone chan struct{} // closed when idle goroutine exits
|
||||
}
|
||||
|
||||
// maxConcurrentRecordings is the maximum number of active (non-stopped)
|
||||
// recordings allowed at once. This prevents resource exhaustion.
|
||||
const maxConcurrentRecordings = 5
|
||||
|
||||
// idleTimeout is the duration of desktop inactivity after which all
|
||||
// active recordings are automatically stopped.
|
||||
const idleTimeout = 10 * time.Minute
|
||||
|
||||
// portableDesktop implements Desktop by shelling out to the
|
||||
// portabledesktop CLI via agentexec.Execer.
|
||||
type portableDesktop struct {
|
||||
logger slog.Logger
|
||||
execer agentexec.Execer
|
||||
scriptBinDir string // coder script bin directory
|
||||
clock quartz.Clock
|
||||
|
||||
mu sync.Mutex
|
||||
session *desktopSession // nil until started
|
||||
binPath string // resolved path to binary, cached
|
||||
closed bool
|
||||
mu sync.Mutex
|
||||
session *desktopSession // nil until started
|
||||
binPath string // resolved path to binary, cached
|
||||
closed bool
|
||||
recordings map[string]*recordingProcess // guarded by mu
|
||||
lastDesktopActionAt atomic.Int64
|
||||
}
|
||||
|
||||
// NewPortableDesktop creates a Desktop backed by the portabledesktop
|
||||
// CLI binary, using execer to spawn child processes. scriptBinDir is
|
||||
// the coder script bin directory checked for the binary.
|
||||
// the coder script bin directory checked for the binary. If clk is
|
||||
// nil, a real clock is used.
|
||||
func NewPortableDesktop(
|
||||
logger slog.Logger,
|
||||
execer agentexec.Execer,
|
||||
scriptBinDir string,
|
||||
clk quartz.Clock,
|
||||
) Desktop {
|
||||
return &portableDesktop{
|
||||
if clk == nil {
|
||||
clk = quartz.NewReal()
|
||||
}
|
||||
pd := &portableDesktop{
|
||||
logger: logger,
|
||||
execer: execer,
|
||||
scriptBinDir: scriptBinDir,
|
||||
clock: clk,
|
||||
recordings: make(map[string]*recordingProcess),
|
||||
}
|
||||
pd.lastDesktopActionAt.Store(clk.Now().UnixNano())
|
||||
return pd
|
||||
}
|
||||
|
||||
// Start launches the desktop session (idempotent).
|
||||
@@ -83,7 +119,7 @@ func (p *portableDesktop) Start(ctx context.Context) (DisplayConfig, error) {
|
||||
defer p.mu.Unlock()
|
||||
|
||||
if p.closed {
|
||||
return DisplayConfig{}, xerrors.New("desktop is closed")
|
||||
return DisplayConfig{}, ErrDesktopClosed
|
||||
}
|
||||
|
||||
if err := p.ensureBinary(ctx); err != nil {
|
||||
@@ -313,23 +349,328 @@ func (p *portableDesktop) CursorPosition(ctx context.Context) (x int, y int, err
|
||||
return result.X, result.Y, nil
|
||||
}
|
||||
|
||||
// Close shuts down the desktop session and cleans up resources.
|
||||
func (p *portableDesktop) Close() error {
|
||||
// StartRecording begins recording the desktop to an MP4 file.
|
||||
// Three-state idempotency: active recordings are no-ops,
|
||||
// completed recordings are discarded and restarted.
|
||||
func (p *portableDesktop) StartRecording(ctx context.Context, recordingID string) error {
|
||||
// Ensure the desktop session is running before acquiring the
|
||||
// recording lock. Start is independently locked and idempotent.
|
||||
if _, err := p.Start(ctx); err != nil {
|
||||
return xerrors.Errorf("ensure desktop session: %w", err)
|
||||
}
|
||||
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
|
||||
if p.closed {
|
||||
return ErrDesktopClosed
|
||||
}
|
||||
|
||||
// Three-state idempotency:
|
||||
// - Active recording → no-op, continue recording.
|
||||
// - Completed recording → discard old file, start fresh.
|
||||
// - Unknown ID → fall through to start a new recording.
|
||||
if rec, ok := p.recordings[recordingID]; ok {
|
||||
if !rec.stopped {
|
||||
select {
|
||||
case <-rec.done:
|
||||
// Process exited unexpectedly; treat as completed
|
||||
// so we fall through to discard the old file and
|
||||
// restart.
|
||||
default:
|
||||
// Active recording - no-op, continue recording.
|
||||
return nil
|
||||
}
|
||||
}
|
||||
// Completed recording - discard old file, start fresh.
|
||||
if err := os.Remove(rec.filePath); err != nil && !os.IsNotExist(err) {
|
||||
p.logger.Warn(ctx, "failed to remove old recording file",
|
||||
slog.F("recording_id", recordingID),
|
||||
slog.F("file_path", rec.filePath),
|
||||
slog.Error(err),
|
||||
)
|
||||
}
|
||||
delete(p.recordings, recordingID)
|
||||
}
|
||||
|
||||
// Check concurrent recording limit.
|
||||
if p.lockedActiveRecordingCount() >= maxConcurrentRecordings {
|
||||
return xerrors.Errorf("too many concurrent recordings (max %d)", maxConcurrentRecordings)
|
||||
}
|
||||
|
||||
// GC sweep: remove stopped recordings with stale files.
|
||||
p.lockedCleanStaleRecordings(ctx)
|
||||
|
||||
if err := p.ensureBinary(ctx); err != nil {
|
||||
return xerrors.Errorf("ensure portabledesktop binary: %w", err)
|
||||
}
|
||||
|
||||
filePath := filepath.Join(os.TempDir(), "coder-recording-"+recordingID+".mp4")
|
||||
|
||||
// Use a background context so the process outlives the HTTP
|
||||
// request that triggered it.
|
||||
procCtx, procCancel := context.WithCancel(context.Background())
|
||||
|
||||
//nolint:gosec // portabledesktop is a trusted binary resolved via ensureBinary.
|
||||
cmd := p.execer.CommandContext(procCtx, p.binPath, "record",
|
||||
// The following options are used to speed up the recording when the desktop is idle.
|
||||
// They were taken out of an example in the portabledesktop repo.
|
||||
// There's likely room for improvement to optimize the values.
|
||||
"--idle-speedup", "20",
|
||||
"--idle-min-duration", "0.35",
|
||||
"--idle-noise-tolerance", "-38dB",
|
||||
filePath)
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
procCancel()
|
||||
return xerrors.Errorf("start recording process: %w", err)
|
||||
}
|
||||
|
||||
rec := &recordingProcess{
|
||||
cmd: cmd,
|
||||
filePath: filePath,
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
go func() {
|
||||
rec.waitErr = cmd.Wait()
|
||||
close(rec.done)
|
||||
// avoid a context resource leak by canceling the context
|
||||
procCancel()
|
||||
}()
|
||||
|
||||
p.recordings[recordingID] = rec
|
||||
|
||||
p.logger.Info(ctx, "started desktop recording",
|
||||
slog.F("recording_id", recordingID),
|
||||
slog.F("file_path", filePath),
|
||||
slog.F("pid", cmd.Process.Pid),
|
||||
)
|
||||
|
||||
// Record activity so a recording started on an already-idle
|
||||
// desktop does not stop immediately.
|
||||
p.lastDesktopActionAt.Store(p.clock.Now().UnixNano())
|
||||
|
||||
// Spawn a per-recording idle goroutine.
|
||||
idleCtx, idleCancel := context.WithCancel(context.Background())
|
||||
rec.idleCancel = idleCancel
|
||||
rec.idleDone = make(chan struct{})
|
||||
go func() {
|
||||
defer close(rec.idleDone)
|
||||
p.monitorRecordingIdle(idleCtx, rec)
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// StopRecording finalizes the recording. Idempotent - safe to call
|
||||
// on an already-stopped recording. Returns a RecordingArtifact
|
||||
// that the caller can stream. The caller must close the Reader
|
||||
// on the returned artifact to avoid leaking file descriptors.
|
||||
func (p *portableDesktop) StopRecording(ctx context.Context, recordingID string) (*RecordingArtifact, error) {
|
||||
p.mu.Lock()
|
||||
rec, ok := p.recordings[recordingID]
|
||||
if !ok {
|
||||
p.mu.Unlock()
|
||||
return nil, ErrUnknownRecording
|
||||
}
|
||||
|
||||
p.lockedStopRecordingProcess(ctx, rec, false)
|
||||
killed := rec.killed
|
||||
p.mu.Unlock()
|
||||
|
||||
p.logger.Info(ctx, "stopped desktop recording",
|
||||
slog.F("recording_id", recordingID),
|
||||
slog.F("file_path", rec.filePath),
|
||||
)
|
||||
|
||||
if killed {
|
||||
return nil, ErrRecordingCorrupted
|
||||
}
|
||||
|
||||
// Open the file and return an artifact. Each call opens a fresh
|
||||
// file descriptor so the caller is insulated from restarts and
|
||||
// desktop close.
|
||||
f, err := os.Open(rec.filePath)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("open recording artifact: %w", err)
|
||||
}
|
||||
info, err := f.Stat()
|
||||
if err != nil {
|
||||
_ = f.Close()
|
||||
return nil, xerrors.Errorf("stat recording artifact: %w", err)
|
||||
}
|
||||
return &RecordingArtifact{
|
||||
Reader: f,
|
||||
Size: info.Size(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// lockedStopRecordingProcess stops a single recording via stopOnce.
|
||||
// It sends SIGINT, waits up to 15 seconds for graceful exit, then
|
||||
// SIGKILLs. When force is true the process is SIGKILLed immediately
|
||||
// without attempting a graceful shutdown. Must be called while p.mu
|
||||
// is held; the lock is held for the full duration so that no
|
||||
// concurrent StopRecording caller can read rec.stopped = true
|
||||
// before the process has finished writing the MP4 file.
|
||||
//
|
||||
//nolint:revive // force flag keeps shared stopOnce/cleanup logic in one place.
|
||||
func (p *portableDesktop) lockedStopRecordingProcess(ctx context.Context, rec *recordingProcess, force bool) {
|
||||
rec.stopOnce.Do(func() {
|
||||
if force {
|
||||
_ = rec.cmd.Process.Kill()
|
||||
rec.killed = true
|
||||
} else {
|
||||
_ = interruptRecordingProcess(rec.cmd.Process)
|
||||
timer := p.clock.NewTimer(15*time.Second, "agentdesktop", "stop_timeout")
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-rec.done:
|
||||
case <-ctx.Done():
|
||||
_ = rec.cmd.Process.Kill()
|
||||
rec.killed = true
|
||||
case <-timer.C:
|
||||
_ = rec.cmd.Process.Kill()
|
||||
rec.killed = true
|
||||
}
|
||||
}
|
||||
rec.stopped = true
|
||||
if rec.idleCancel != nil {
|
||||
rec.idleCancel()
|
||||
}
|
||||
})
|
||||
// NOTE: We intentionally do not wait on rec.done here.
|
||||
// If goleak is added to this package's tests, this may
|
||||
// need revisiting to avoid flakes.
|
||||
}
|
||||
|
||||
// lockedActiveRecordingCount returns the number of recordings that
|
||||
// are still actively running. Must be called while p.mu is held.
|
||||
// The max concurrency is low (maxConcurrentRecordings = 5), so a
|
||||
// full scan is cheap and avoids maintaining a separate counter.
|
||||
func (p *portableDesktop) lockedActiveRecordingCount() int {
|
||||
active := 0
|
||||
for _, rec := range p.recordings {
|
||||
if rec.stopped {
|
||||
continue
|
||||
}
|
||||
select {
|
||||
case <-rec.done:
|
||||
default:
|
||||
active++
|
||||
}
|
||||
}
|
||||
return active
|
||||
}
|
||||
|
||||
// lockedCleanStaleRecordings removes stopped recordings whose temp
|
||||
// files are older than one hour. Must be called while p.mu is held.
|
||||
func (p *portableDesktop) lockedCleanStaleRecordings(ctx context.Context) {
|
||||
for id, rec := range p.recordings {
|
||||
if !rec.stopped {
|
||||
continue
|
||||
}
|
||||
info, err := os.Stat(rec.filePath)
|
||||
if err != nil {
|
||||
// File already removed or inaccessible; drop entry.
|
||||
delete(p.recordings, id)
|
||||
continue
|
||||
}
|
||||
if p.clock.Since(info.ModTime()) > time.Hour {
|
||||
if err := os.Remove(rec.filePath); err != nil && !os.IsNotExist(err) {
|
||||
p.logger.Warn(ctx, "failed to remove stale recording file",
|
||||
slog.F("recording_id", id),
|
||||
slog.F("file_path", rec.filePath),
|
||||
slog.Error(err),
|
||||
)
|
||||
}
|
||||
delete(p.recordings, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Close shuts down the desktop session and cleans up resources.
|
||||
func (p *portableDesktop) Close() error {
|
||||
p.mu.Lock()
|
||||
p.closed = true
|
||||
if p.session != nil {
|
||||
p.session.cancel()
|
||||
// Xvnc is a child process — killing it cleans up the X
|
||||
// session.
|
||||
_ = p.session.cmd.Process.Kill()
|
||||
_ = p.session.cmd.Wait()
|
||||
p.session = nil
|
||||
|
||||
// Force-kill all active recordings. The stopOnce inside
|
||||
// lockedStopRecordingProcess makes this safe for
|
||||
// already-stopped recordings.
|
||||
for _, rec := range p.recordings {
|
||||
p.lockedStopRecordingProcess(context.Background(), rec, true)
|
||||
}
|
||||
|
||||
// Snapshot recording file paths and idle goroutine channels
|
||||
// for cleanup, then clear the map.
|
||||
type recEntry struct {
|
||||
id string
|
||||
filePath string
|
||||
idleDone chan struct{}
|
||||
}
|
||||
var allRecs []recEntry
|
||||
for id, rec := range p.recordings {
|
||||
allRecs = append(allRecs, recEntry{id: id, filePath: rec.filePath, idleDone: rec.idleDone})
|
||||
delete(p.recordings, id)
|
||||
}
|
||||
session := p.session
|
||||
p.session = nil
|
||||
p.mu.Unlock()
|
||||
|
||||
// Wait for all per-recording idle goroutines to exit.
|
||||
for _, entry := range allRecs {
|
||||
if entry.idleDone != nil {
|
||||
<-entry.idleDone
|
||||
}
|
||||
}
|
||||
|
||||
// Remove all recording files and wait for the session to
|
||||
// exit with a timeout so a slow filesystem or hung process
|
||||
// cannot block agent shutdown indefinitely.
|
||||
cleanupDone := make(chan struct{})
|
||||
go func() {
|
||||
defer close(cleanupDone)
|
||||
for _, entry := range allRecs {
|
||||
if err := os.Remove(entry.filePath); err != nil && !os.IsNotExist(err) {
|
||||
p.logger.Warn(context.Background(), "failed to remove recording file on close",
|
||||
slog.F("recording_id", entry.id),
|
||||
slog.F("file_path", entry.filePath),
|
||||
slog.Error(err),
|
||||
)
|
||||
}
|
||||
}
|
||||
if session != nil {
|
||||
session.cancel()
|
||||
if err := session.cmd.Process.Kill(); err != nil {
|
||||
p.logger.Warn(context.Background(), "failed to kill portabledesktop process",
|
||||
slog.Error(err),
|
||||
)
|
||||
}
|
||||
if err := session.cmd.Wait(); err != nil {
|
||||
var exitErr *exec.ExitError
|
||||
if !errors.As(err, &exitErr) {
|
||||
p.logger.Warn(context.Background(), "portabledesktop process exited with error",
|
||||
slog.Error(err),
|
||||
)
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
timer := p.clock.NewTimer(15*time.Second, "agentdesktop", "close_cleanup_timeout")
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-cleanupDone:
|
||||
case <-timer.C:
|
||||
p.logger.Warn(context.Background(), "timed out waiting for close cleanup")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// RecordActivity marks the desktop as having received user
|
||||
// interaction, resetting the idle-recording timer.
|
||||
func (p *portableDesktop) RecordActivity() {
|
||||
p.lastDesktopActionAt.Store(p.clock.Now().UnixNano())
|
||||
}
|
||||
|
||||
// runCmd executes a portabledesktop subcommand and returns combined
|
||||
// output. The caller must have previously called ensureBinary.
|
||||
func (p *portableDesktop) runCmd(ctx context.Context, args ...string) (string, error) {
|
||||
@@ -397,3 +738,31 @@ func (p *portableDesktop) ensureBinary(ctx context.Context) error {
|
||||
|
||||
return xerrors.New("portabledesktop binary not found in PATH or script bin directory")
|
||||
}
|
||||
|
||||
// monitorRecordingIdle watches for desktop inactivity and stops the
|
||||
// given recording when the idle timeout is reached.
|
||||
func (p *portableDesktop) monitorRecordingIdle(ctx context.Context, rec *recordingProcess) {
|
||||
timer := p.clock.NewTimer(idleTimeout, "agentdesktop", "recording_idle")
|
||||
defer timer.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-timer.C:
|
||||
lastNano := p.lastDesktopActionAt.Load()
|
||||
lastAction := time.Unix(0, lastNano)
|
||||
elapsed := p.clock.Since(lastAction)
|
||||
if elapsed >= idleTimeout {
|
||||
p.mu.Lock()
|
||||
p.lockedStopRecordingProcess(context.Background(), rec, false)
|
||||
p.mu.Unlock()
|
||||
return
|
||||
}
|
||||
// Activity happened; reset with remaining budget.
|
||||
timer.Reset(idleTimeout-elapsed, "agentdesktop", "recording_idle")
|
||||
case <-rec.done:
|
||||
return
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,13 +9,17 @@ import (
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
"github.com/coder/coder/v2/agent/agentexec"
|
||||
"github.com/coder/coder/v2/pty"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
// recordedExecer implements agentexec.Execer by recording every
|
||||
@@ -86,6 +90,7 @@ func TestPortableDesktop_Start_ParsesOutput(t *testing.T) {
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
binPath: "portabledesktop", // pre-set so ensureBinary is a no-op
|
||||
clock: quartz.NewReal(),
|
||||
}
|
||||
|
||||
ctx := t.Context()
|
||||
@@ -117,6 +122,7 @@ func TestPortableDesktop_Start_Idempotent(t *testing.T) {
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
binPath: "portabledesktop",
|
||||
clock: quartz.NewReal(),
|
||||
}
|
||||
|
||||
ctx := t.Context()
|
||||
@@ -159,6 +165,7 @@ func TestPortableDesktop_Screenshot(t *testing.T) {
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
binPath: "portabledesktop",
|
||||
clock: quartz.NewReal(),
|
||||
}
|
||||
|
||||
ctx := t.Context()
|
||||
@@ -184,6 +191,7 @@ func TestPortableDesktop_Screenshot_WithTargetDimensions(t *testing.T) {
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
binPath: "portabledesktop",
|
||||
clock: quartz.NewReal(),
|
||||
}
|
||||
|
||||
ctx := t.Context()
|
||||
@@ -282,6 +290,7 @@ func TestPortableDesktop_MouseMethods(t *testing.T) {
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
binPath: "portabledesktop",
|
||||
clock: quartz.NewReal(),
|
||||
}
|
||||
|
||||
err := tt.invoke(t.Context(), pd)
|
||||
@@ -289,7 +298,6 @@ func TestPortableDesktop_MouseMethods(t *testing.T) {
|
||||
|
||||
cmds := rec.allCommands()
|
||||
require.NotEmpty(t, cmds, "expected at least one command")
|
||||
|
||||
// Find at least one recorded command that contains
|
||||
// all expected argument substrings.
|
||||
found := false
|
||||
@@ -367,6 +375,7 @@ func TestPortableDesktop_KeyboardMethods(t *testing.T) {
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
binPath: "portabledesktop",
|
||||
clock: quartz.NewReal(),
|
||||
}
|
||||
|
||||
err := tt.invoke(t.Context(), pd)
|
||||
@@ -423,6 +432,7 @@ func TestPortableDesktop_Close(t *testing.T) {
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
binPath: "portabledesktop",
|
||||
clock: quartz.NewReal(),
|
||||
}
|
||||
|
||||
ctx := t.Context()
|
||||
@@ -445,7 +455,7 @@ func TestPortableDesktop_Close(t *testing.T) {
|
||||
// Subsequent Start must fail.
|
||||
_, err = pd.Start(ctx)
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "desktop is closed")
|
||||
assert.Contains(t, err.Error(), "desktop closed")
|
||||
}
|
||||
|
||||
// --- ensureBinary tests ---
|
||||
@@ -539,7 +549,410 @@ func TestEnsureBinary_NotFound(t *testing.T) {
|
||||
assert.Contains(t, err.Error(), "not found")
|
||||
}
|
||||
|
||||
func TestPortableDesktop_StartRecording(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
rec := &recordedExecer{
|
||||
scripts: map[string]string{
|
||||
"record": `trap 'exit 0' INT; sleep 120 & wait`,
|
||||
"up": `printf '{"vncPort":5901,"geometry":"1920x1080"}\n' && sleep 120`,
|
||||
},
|
||||
}
|
||||
|
||||
clk := quartz.NewReal()
|
||||
pd := &portableDesktop{
|
||||
logger: logger,
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
clock: clk,
|
||||
binPath: "portabledesktop",
|
||||
recordings: make(map[string]*recordingProcess),
|
||||
}
|
||||
pd.lastDesktopActionAt.Store(clk.Now().UnixNano())
|
||||
|
||||
ctx := t.Context()
|
||||
recID := uuid.New().String()
|
||||
err := pd.StartRecording(ctx, recID)
|
||||
require.NoError(t, err)
|
||||
|
||||
cmds := rec.allCommands()
|
||||
require.NotEmpty(t, cmds)
|
||||
// Find the record command (not the up command).
|
||||
found := false
|
||||
for _, cmd := range cmds {
|
||||
joined := strings.Join(cmd, " ")
|
||||
if strings.Contains(joined, "record") && strings.Contains(joined, "coder-recording-"+recID) {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
assert.True(t, found, "expected a record command with the recording ID")
|
||||
|
||||
require.NoError(t, pd.Close())
|
||||
}
|
||||
|
||||
func TestPortableDesktop_StartRecording_ConcurrentLimit(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
rec := &recordedExecer{
|
||||
scripts: map[string]string{
|
||||
"record": `trap 'exit 0' INT; sleep 120 & wait`,
|
||||
"up": `printf '{"vncPort":5901,"geometry":"1920x1080"}\n' && sleep 120`,
|
||||
},
|
||||
}
|
||||
|
||||
clk := quartz.NewReal()
|
||||
pd := &portableDesktop{
|
||||
logger: logger,
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
clock: clk,
|
||||
binPath: "portabledesktop",
|
||||
recordings: make(map[string]*recordingProcess),
|
||||
}
|
||||
pd.lastDesktopActionAt.Store(clk.Now().UnixNano())
|
||||
|
||||
ctx := t.Context()
|
||||
|
||||
for i := range maxConcurrentRecordings {
|
||||
err := pd.StartRecording(ctx, uuid.New().String())
|
||||
require.NoError(t, err, "recording %d should succeed", i)
|
||||
}
|
||||
|
||||
err := pd.StartRecording(ctx, uuid.New().String())
|
||||
require.Error(t, err)
|
||||
assert.Contains(t, err.Error(), "too many concurrent recordings")
|
||||
|
||||
require.NoError(t, pd.Close())
|
||||
}
|
||||
|
||||
func TestPortableDesktop_StopRecording_ReturnsArtifact(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
rec := &recordedExecer{
|
||||
scripts: map[string]string{
|
||||
"record": `trap 'exit 0' INT; sleep 120 & wait`,
|
||||
"up": `printf '{"vncPort":5901,"geometry":"1920x1080"}\n' && sleep 120`,
|
||||
},
|
||||
}
|
||||
|
||||
clk := quartz.NewReal()
|
||||
pd := &portableDesktop{
|
||||
logger: logger,
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
clock: clk,
|
||||
binPath: "portabledesktop",
|
||||
recordings: make(map[string]*recordingProcess),
|
||||
}
|
||||
pd.lastDesktopActionAt.Store(clk.Now().UnixNano())
|
||||
|
||||
ctx := t.Context()
|
||||
recID := uuid.New().String()
|
||||
err := pd.StartRecording(ctx, recID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Write a dummy MP4 file at the expected path so StopRecording
|
||||
// can open it as an artifact.
|
||||
filePath := filepath.Join(os.TempDir(), "coder-recording-"+recID+".mp4")
|
||||
require.NoError(t, os.WriteFile(filePath, []byte("fake-mp4-data"), 0o600))
|
||||
t.Cleanup(func() { _ = os.Remove(filePath) })
|
||||
|
||||
artifact, err := pd.StopRecording(ctx, recID)
|
||||
require.NoError(t, err)
|
||||
defer artifact.Reader.Close()
|
||||
assert.Equal(t, int64(len("fake-mp4-data")), artifact.Size)
|
||||
|
||||
require.NoError(t, pd.Close())
|
||||
}
|
||||
|
||||
func TestPortableDesktop_StopRecording_UnknownID(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
rec := &recordedExecer{
|
||||
scripts: map[string]string{
|
||||
"record": `trap 'exit 0' INT; sleep 120 & wait`,
|
||||
},
|
||||
}
|
||||
|
||||
clk := quartz.NewReal()
|
||||
pd := &portableDesktop{
|
||||
logger: logger,
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
clock: clk,
|
||||
binPath: "portabledesktop",
|
||||
recordings: make(map[string]*recordingProcess),
|
||||
}
|
||||
pd.lastDesktopActionAt.Store(clk.Now().UnixNano())
|
||||
|
||||
ctx := t.Context()
|
||||
_, err := pd.StopRecording(ctx, uuid.New().String())
|
||||
require.ErrorIs(t, err, ErrUnknownRecording)
|
||||
|
||||
require.NoError(t, pd.Close())
|
||||
}
|
||||
|
||||
// Ensure that portableDesktop satisfies the Desktop interface at
|
||||
// compile time. This uses the unexported type so it lives in the
|
||||
// internal test package.
|
||||
var _ Desktop = (*portableDesktop)(nil)
|
||||
|
||||
func TestPortableDesktop_IdleTimeout_StopsRecordings(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
rec := &recordedExecer{
|
||||
scripts: map[string]string{
|
||||
"record": `trap 'exit 0' INT; sleep 120 & wait`,
|
||||
"up": `printf '{"vncPort":5901,"geometry":"1920x1080"}\n' && sleep 120`,
|
||||
},
|
||||
}
|
||||
|
||||
clk := quartz.NewMock(t)
|
||||
pd := &portableDesktop{
|
||||
logger: logger,
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
clock: clk,
|
||||
binPath: "portabledesktop",
|
||||
recordings: make(map[string]*recordingProcess),
|
||||
}
|
||||
pd.lastDesktopActionAt.Store(clk.Now().UnixNano())
|
||||
|
||||
ctx := t.Context()
|
||||
recID := uuid.New().String()
|
||||
|
||||
// Install the trap before StartRecording so it is guaranteed
|
||||
// to catch the idle monitor's NewTimer call regardless of
|
||||
// goroutine scheduling.
|
||||
trap := clk.Trap().NewTimer("agentdesktop", "recording_idle")
|
||||
|
||||
err := pd.StartRecording(ctx, recID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify recording is active.
|
||||
pd.mu.Lock()
|
||||
require.False(t, pd.recordings[recID].stopped)
|
||||
pd.mu.Unlock()
|
||||
|
||||
// Wait for the idle monitor timer to be created and release
|
||||
// it so the monitor enters its select loop.
|
||||
trap.MustWait(ctx).MustRelease(ctx)
|
||||
trap.Close()
|
||||
|
||||
// The stop-all path calls lockedStopRecordingProcess which
|
||||
// creates a per-recording 15s stop_timeout timer.
|
||||
stopTrap := clk.Trap().NewTimer("agentdesktop", "stop_timeout")
|
||||
|
||||
// Advance past idle timeout to trigger the stop-all.
|
||||
clk.Advance(idleTimeout)
|
||||
|
||||
// Wait for the stop timer to be created, then release it.
|
||||
stopTrap.MustWait(ctx).MustRelease(ctx)
|
||||
stopTrap.Close()
|
||||
|
||||
// The recording process should now be stopped.
|
||||
require.Eventually(t, func() bool {
|
||||
pd.mu.Lock()
|
||||
defer pd.mu.Unlock()
|
||||
rec, ok := pd.recordings[recID]
|
||||
return ok && rec.stopped
|
||||
}, testutil.WaitShort, testutil.IntervalFast)
|
||||
|
||||
require.NoError(t, pd.Close())
|
||||
}
|
||||
|
||||
func TestPortableDesktop_IdleTimeout_ActivityResetsTimer(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
rec := &recordedExecer{
|
||||
scripts: map[string]string{
|
||||
"record": `trap 'exit 0' INT; sleep 120 & wait`,
|
||||
"up": `printf '{"vncPort":5901,"geometry":"1920x1080"}\n' && sleep 120`,
|
||||
},
|
||||
}
|
||||
|
||||
clk := quartz.NewMock(t)
|
||||
pd := &portableDesktop{
|
||||
logger: logger,
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
clock: clk,
|
||||
binPath: "portabledesktop",
|
||||
recordings: make(map[string]*recordingProcess),
|
||||
}
|
||||
pd.lastDesktopActionAt.Store(clk.Now().UnixNano())
|
||||
|
||||
ctx := t.Context()
|
||||
recID := uuid.New().String()
|
||||
|
||||
// Install the trap before StartRecording so it is guaranteed
|
||||
// to catch the idle monitor's NewTimer call regardless of
|
||||
// goroutine scheduling.
|
||||
trap := clk.Trap().NewTimer("agentdesktop", "recording_idle")
|
||||
|
||||
err := pd.StartRecording(ctx, recID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Wait for the idle monitor timer to be created.
|
||||
trap.MustWait(ctx).MustRelease(ctx)
|
||||
trap.Close()
|
||||
|
||||
// Advance most of the way but not past the timeout.
|
||||
clk.Advance(idleTimeout - time.Minute)
|
||||
|
||||
// Record activity to reset the timer.
|
||||
pd.RecordActivity()
|
||||
|
||||
// Trap the Reset call that the idle monitor makes when it
|
||||
// sees recent activity.
|
||||
resetTrap := clk.Trap().TimerReset("agentdesktop", "recording_idle")
|
||||
|
||||
// Advance past the original idle timeout deadline. The
|
||||
// monitor should see the recent activity and reset instead
|
||||
// of stopping.
|
||||
clk.Advance(time.Minute)
|
||||
|
||||
resetTrap.MustWait(ctx).MustRelease(ctx)
|
||||
resetTrap.Close()
|
||||
|
||||
// Recording should still be active because activity was
|
||||
// recorded.
|
||||
pd.mu.Lock()
|
||||
require.False(t, pd.recordings[recID].stopped)
|
||||
pd.mu.Unlock()
|
||||
|
||||
require.NoError(t, pd.Close())
|
||||
}
|
||||
|
||||
func TestPortableDesktop_IdleTimeout_MultipleRecordings(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
rec := &recordedExecer{
|
||||
scripts: map[string]string{
|
||||
"record": `trap 'exit 0' INT; sleep 120 & wait`,
|
||||
"up": `printf '{"vncPort":5901,"geometry":"1920x1080"}\n' && sleep 120`,
|
||||
},
|
||||
}
|
||||
|
||||
clk := quartz.NewMock(t)
|
||||
pd := &portableDesktop{
|
||||
logger: logger,
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
clock: clk,
|
||||
binPath: "portabledesktop",
|
||||
recordings: make(map[string]*recordingProcess),
|
||||
}
|
||||
pd.lastDesktopActionAt.Store(clk.Now().UnixNano())
|
||||
|
||||
ctx := t.Context()
|
||||
recID1 := uuid.New().String()
|
||||
recID2 := uuid.New().String()
|
||||
|
||||
// Trap idle timer creation for both recordings.
|
||||
trap := clk.Trap().NewTimer("agentdesktop", "recording_idle")
|
||||
|
||||
err := pd.StartRecording(ctx, recID1)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Wait for first recording's idle timer.
|
||||
trap.MustWait(ctx).MustRelease(ctx)
|
||||
|
||||
err = pd.StartRecording(ctx, recID2)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Wait for second recording's idle timer.
|
||||
trap.MustWait(ctx).MustRelease(ctx)
|
||||
trap.Close()
|
||||
|
||||
// Trap the stop timers that will be created when idle fires.
|
||||
stopTrap := clk.Trap().NewTimer("agentdesktop", "stop_timeout")
|
||||
|
||||
// Advance past idle timeout.
|
||||
clk.Advance(idleTimeout)
|
||||
|
||||
// Wait for both stop timers.
|
||||
stopTrap.MustWait(ctx).MustRelease(ctx)
|
||||
stopTrap.MustWait(ctx).MustRelease(ctx)
|
||||
stopTrap.Close()
|
||||
|
||||
// Both recordings should be stopped.
|
||||
require.Eventually(t, func() bool {
|
||||
pd.mu.Lock()
|
||||
defer pd.mu.Unlock()
|
||||
r1, ok1 := pd.recordings[recID1]
|
||||
r2, ok2 := pd.recordings[recID2]
|
||||
return ok1 && r1.stopped && ok2 && r2.stopped
|
||||
}, testutil.WaitShort, testutil.IntervalFast)
|
||||
|
||||
require.NoError(t, pd.Close())
|
||||
}
|
||||
|
||||
func TestPortableDesktop_StartRecording_ReturnsErrDesktopClosed(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
rec := &recordedExecer{
|
||||
scripts: map[string]string{
|
||||
"up": `printf '{"vncPort":5901,"geometry":"1920x1080"}\n' && sleep 120`,
|
||||
},
|
||||
}
|
||||
|
||||
clk := quartz.NewReal()
|
||||
pd := &portableDesktop{
|
||||
logger: logger,
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
clock: clk,
|
||||
binPath: "portabledesktop",
|
||||
recordings: make(map[string]*recordingProcess),
|
||||
}
|
||||
pd.lastDesktopActionAt.Store(clk.Now().UnixNano())
|
||||
|
||||
// Start and close the desktop so it's in the closed state.
|
||||
ctx := t.Context()
|
||||
_, err := pd.Start(ctx)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, pd.Close())
|
||||
|
||||
// StartRecording should now return ErrDesktopClosed.
|
||||
err = pd.StartRecording(ctx, uuid.New().String())
|
||||
require.ErrorIs(t, err, ErrDesktopClosed)
|
||||
}
|
||||
|
||||
func TestPortableDesktop_Start_ReturnsErrDesktopClosed(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slogtest.Make(t, nil)
|
||||
rec := &recordedExecer{
|
||||
scripts: map[string]string{
|
||||
"up": `printf '{"vncPort":5901,"geometry":"1920x1080"}\n' && sleep 120`,
|
||||
},
|
||||
}
|
||||
|
||||
pd := &portableDesktop{
|
||||
logger: logger,
|
||||
execer: rec,
|
||||
scriptBinDir: t.TempDir(),
|
||||
clock: quartz.NewReal(),
|
||||
binPath: "portabledesktop",
|
||||
recordings: make(map[string]*recordingProcess),
|
||||
}
|
||||
pd.lastDesktopActionAt.Store(pd.clock.Now().UnixNano())
|
||||
|
||||
ctx := t.Context()
|
||||
_, err := pd.Start(ctx)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, pd.Close())
|
||||
|
||||
_, err = pd.Start(ctx)
|
||||
require.ErrorIs(t, err, ErrDesktopClosed)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
//go:build !windows
|
||||
|
||||
package agentdesktop
|
||||
|
||||
import "os"
|
||||
|
||||
// interruptRecordingProcess sends a SIGINT to the recording process
|
||||
// for graceful shutdown. On Unix, os.Interrupt is delivered as
|
||||
// SIGINT which lets the recorder finalize the MP4 container.
|
||||
func interruptRecordingProcess(p *os.Process) error {
|
||||
return p.Signal(os.Interrupt)
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
package agentdesktop
|
||||
|
||||
import "os"
|
||||
|
||||
// interruptRecordingProcess kills the recording process directly
|
||||
// because os.Process.Signal(os.Interrupt) is not supported on
|
||||
// Windows and returns an error without delivering a signal.
|
||||
func interruptRecordingProcess(p *os.Process) error {
|
||||
return p.Kill()
|
||||
}
|
||||
@@ -71,6 +71,13 @@ const (
|
||||
// events cached per chat for same-replica stream catch-up.
|
||||
maxDurableMessageCacheSize = 256
|
||||
|
||||
// maxConcurrentRecordingUploads caps the number of recording
|
||||
// stop-and-store operations that can run concurrently. Each
|
||||
// slot buffers up to MaxRecordingSize (100 MB) in memory, so
|
||||
// this value implicitly bounds memory to roughly
|
||||
// maxConcurrentRecordingUploads * 100 MB.
|
||||
maxConcurrentRecordingUploads = 25
|
||||
|
||||
// staleRecoveryIntervalDivisor determines how often the stale
|
||||
// recovery loop runs relative to the stale threshold. A value
|
||||
// of 5 means recovery runs at 1/5 of the stale-after duration.
|
||||
@@ -129,6 +136,7 @@ type Server struct {
|
||||
|
||||
usageTracker *workspacestats.UsageTracker
|
||||
clock quartz.Clock
|
||||
recordingSem chan struct{}
|
||||
|
||||
// Configuration
|
||||
pendingChatAcquireInterval time.Duration
|
||||
@@ -2372,6 +2380,7 @@ func New(cfg Config) *Server {
|
||||
chatHeartbeatInterval: chatHeartbeatInterval,
|
||||
usageTracker: cfg.UsageTracker,
|
||||
clock: clk,
|
||||
recordingSem: make(chan struct{}, maxConcurrentRecordingUploads),
|
||||
wakeCh: make(chan struct{}, 1),
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
package chatd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbauthz"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
)
|
||||
|
||||
// stopAndStoreRecording stops the desktop recording, downloads the
|
||||
// MP4, and stores it in chat_files. Only called when the subagent
|
||||
// completed successfully. Returns the file ID on success, empty
|
||||
// string on any failure. All errors are logged but not propagated
|
||||
// — recording is best-effort.
|
||||
func (p *Server) stopAndStoreRecording(
|
||||
ctx context.Context,
|
||||
conn workspacesdk.AgentConn,
|
||||
recordingID string,
|
||||
ownerID uuid.UUID,
|
||||
workspaceID uuid.NullUUID,
|
||||
) string {
|
||||
select {
|
||||
case p.recordingSem <- struct{}{}:
|
||||
defer func() { <-p.recordingSem }()
|
||||
case <-ctx.Done():
|
||||
p.logger.Warn(ctx, "context canceled waiting for recording semaphore", slog.Error(ctx.Err()))
|
||||
return ""
|
||||
}
|
||||
|
||||
body, err := conn.StopDesktopRecording(ctx,
|
||||
workspacesdk.StopDesktopRecordingRequest{RecordingID: recordingID})
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "failed to stop desktop recording",
|
||||
slog.Error(err))
|
||||
return ""
|
||||
}
|
||||
type readResult struct {
|
||||
data []byte
|
||||
err error
|
||||
}
|
||||
ch := make(chan readResult, 1)
|
||||
go func() {
|
||||
data, err := io.ReadAll(io.LimitReader(body, workspacesdk.MaxRecordingSize+1))
|
||||
ch <- readResult{data, err}
|
||||
}()
|
||||
|
||||
var data []byte
|
||||
select {
|
||||
case res := <-ch:
|
||||
body.Close()
|
||||
data = res.data
|
||||
if res.err != nil {
|
||||
p.logger.Warn(ctx, "failed to read recording data", slog.Error(res.err))
|
||||
return ""
|
||||
}
|
||||
case <-ctx.Done():
|
||||
body.Close()
|
||||
p.logger.Warn(ctx, "context canceled while reading recording data", slog.Error(ctx.Err()))
|
||||
return ""
|
||||
}
|
||||
if len(data) > workspacesdk.MaxRecordingSize {
|
||||
p.logger.Warn(ctx, "recording data exceeds maximum size, skipping store",
|
||||
slog.F("size", len(data)),
|
||||
slog.F("max_size", workspacesdk.MaxRecordingSize))
|
||||
return ""
|
||||
}
|
||||
if len(data) == 0 {
|
||||
p.logger.Warn(ctx, "recording data is empty, skipping store")
|
||||
return ""
|
||||
}
|
||||
|
||||
if !workspaceID.Valid {
|
||||
p.logger.Warn(ctx, "chat has no workspace, cannot store recording")
|
||||
return ""
|
||||
}
|
||||
|
||||
// The chatd actor is used here because the recording is stored on
|
||||
// behalf of the chat system, not a specific user request.
|
||||
//nolint:gocritic // AsChatd is required to read the workspace for org lookup.
|
||||
ws, err := p.db.GetWorkspaceByID(dbauthz.AsChatd(ctx), workspaceID.UUID)
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "failed to resolve workspace for recording",
|
||||
slog.Error(err))
|
||||
return ""
|
||||
}
|
||||
|
||||
//nolint:gocritic // AsChatd is required to insert chat files from the recording pipeline.
|
||||
row, err := p.db.InsertChatFile(dbauthz.AsChatd(ctx), database.InsertChatFileParams{
|
||||
OwnerID: ownerID,
|
||||
OrganizationID: ws.OrganizationID,
|
||||
Name: fmt.Sprintf("recording-%s.mp4", p.clock.Now().UTC().Format("2006-01-02T15-04-05Z")),
|
||||
Mimetype: "video/mp4",
|
||||
Data: data,
|
||||
})
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "failed to store recording in database",
|
||||
slog.Error(err))
|
||||
return ""
|
||||
}
|
||||
return row.ID.String()
|
||||
}
|
||||
@@ -0,0 +1,514 @@
|
||||
package chatd
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"charm.land/fantasy"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
// zeroReader is an io.Reader that produces zero-valued bytes
|
||||
// without allocating large buffers.
|
||||
type zeroReader struct{}
|
||||
|
||||
func (zeroReader) Read(p []byte) (int, error) {
|
||||
clear(p)
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
// createComputerUseParentChild creates a parent chat and a
|
||||
// computer_use child chat bound to the given workspace/agent.
|
||||
// Both chats are inserted directly via DB to avoid triggering
|
||||
// background processing (which would try to call the LLM and
|
||||
// use the agent connection mock).
|
||||
func createComputerUseParentChild(
|
||||
ctx context.Context,
|
||||
t *testing.T,
|
||||
server *Server,
|
||||
user database.User,
|
||||
model database.ChatModelConfig,
|
||||
workspace database.WorkspaceTable,
|
||||
agent database.WorkspaceAgent,
|
||||
parentTitle, childTitle string,
|
||||
) (parent, child database.Chat) {
|
||||
t.Helper()
|
||||
|
||||
// Insert the parent chat directly via DB to avoid triggering
|
||||
// the server's background processing.
|
||||
parent, err := server.db.InsertChat(ctx, database.InsertChatParams{
|
||||
OwnerID: user.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
AgentID: uuid.NullUUID{UUID: agent.ID, Valid: true},
|
||||
LastModelConfigID: model.ID,
|
||||
Title: parentTitle,
|
||||
Status: database.ChatStatusPending,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Insert the child chat directly via DB to avoid triggering
|
||||
// the server's background processing (which would try to run
|
||||
// the chat without an LLM and get stuck).
|
||||
child, err = server.db.InsertChat(ctx, database.InsertChatParams{
|
||||
OwnerID: user.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
AgentID: uuid.NullUUID{UUID: agent.ID, Valid: true},
|
||||
ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true},
|
||||
RootChatID: uuid.NullUUID{UUID: parent.ID, Valid: true},
|
||||
LastModelConfigID: model.ID,
|
||||
Title: childTitle,
|
||||
Mode: database.NullChatMode{ChatMode: database.ChatModeComputerUse, Valid: true},
|
||||
Status: database.ChatStatusPending,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
return parent, child
|
||||
}
|
||||
|
||||
// invokeWaitAgentTool builds the wait_agent tool from the server and
|
||||
// invokes it with the given child chat ID and timeout.
|
||||
func invokeWaitAgentTool(
|
||||
ctx context.Context,
|
||||
t *testing.T,
|
||||
server *Server,
|
||||
db database.Store,
|
||||
parentID uuid.UUID,
|
||||
childID uuid.UUID,
|
||||
timeoutSeconds int,
|
||||
) (fantasy.ToolResponse, error) {
|
||||
t.Helper()
|
||||
|
||||
// Re-fetch the parent so LastModelConfigID is populated.
|
||||
parentChat, err := db.GetChatByID(ctx, parentID)
|
||||
require.NoError(t, err)
|
||||
|
||||
tools := server.subagentTools(ctx, func() database.Chat { return parentChat })
|
||||
tool := findToolByName(tools, "wait_agent")
|
||||
require.NotNil(t, tool, "wait_agent tool must be present")
|
||||
|
||||
argsJSON, err := json.Marshal(map[string]any{
|
||||
"chat_id": childID.String(),
|
||||
"timeout_seconds": timeoutSeconds,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
return tool.Run(ctx, fantasy.ToolCall{
|
||||
ID: "test-call",
|
||||
Name: "wait_agent",
|
||||
Input: string(argsJSON),
|
||||
})
|
||||
}
|
||||
|
||||
// TestWaitAgentComputerUseRecording verifies the happy-path recording
|
||||
// flow: for a computer_use child chat that completes successfully,
|
||||
// the recording is stopped, the MP4 is stored in chat_files, and the
|
||||
// file ID is returned.
|
||||
func TestWaitAgentComputerUseRecording(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := chatdTestContext(t)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
|
||||
user, model := seedInternalChatDeps(ctx, t, db)
|
||||
workspace, _, agent := seedWorkspaceBinding(t, db, user.ID)
|
||||
|
||||
// Create the server WITHOUT agentConnFn so the background
|
||||
// processing of the parent chat doesn't use the mock.
|
||||
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
||||
|
||||
parent, child := createComputerUseParentChild(
|
||||
ctx, t, server, user, model, workspace, agent,
|
||||
"parent-recording", "computer-use-child",
|
||||
)
|
||||
|
||||
// Wait for background processing triggered by CreateChat to
|
||||
// settle before setting up the mock agent connection.
|
||||
server.inflight.Wait()
|
||||
|
||||
// Now wire up the mock agent connection.
|
||||
server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
||||
require.Equal(t, agent.ID, agentID)
|
||||
return mockConn, func() {}, nil
|
||||
}
|
||||
|
||||
// Add an assistant message so the report is extracted.
|
||||
insertAssistantMessage(ctx, t, db, child.ID, model.ID, "I opened Firefox.")
|
||||
|
||||
// Set child to waiting (terminal success state).
|
||||
setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
|
||||
|
||||
// Set up mock expectations for start and stop.
|
||||
fakeMp4 := []byte("fake-mp4-data-for-recording-test")
|
||||
|
||||
mockConn.EXPECT().
|
||||
StartDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(_ context.Context, req workspacesdk.StartDesktopRecordingRequest) error {
|
||||
require.NotEmpty(t, req.RecordingID, "recording ID should be non-empty")
|
||||
return nil
|
||||
}).
|
||||
Times(1)
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(io.NopCloser(bytes.NewReader(fakeMp4)), nil).
|
||||
Times(1)
|
||||
|
||||
// Invoke wait_agent via the tool closure.
|
||||
resp, err := invokeWaitAgentTool(ctx, t, server, db, parent.ID, child.ID, 5)
|
||||
require.NoError(t, err)
|
||||
require.False(t, resp.IsError, "expected successful response, got: %s", resp.Content)
|
||||
|
||||
// Parse the response JSON and check for recording_file_id.
|
||||
var result map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
||||
storedFileID, ok := result["recording_file_id"].(string)
|
||||
require.True(t, ok, "recording_file_id must be present in response")
|
||||
require.NotEmpty(t, storedFileID)
|
||||
|
||||
// Verify the file was inserted into the database.
|
||||
fileUUID, err := uuid.Parse(storedFileID)
|
||||
require.NoError(t, err)
|
||||
|
||||
chatFile, err := db.GetChatFileByID(ctx, fileUUID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "video/mp4", chatFile.Mimetype)
|
||||
assert.True(t, strings.HasPrefix(chatFile.Name, "recording-"),
|
||||
"expected name to start with 'recording-', got: %s", chatFile.Name)
|
||||
assert.Equal(t, user.ID, chatFile.OwnerID)
|
||||
assert.Equal(t, fakeMp4, chatFile.Data)
|
||||
}
|
||||
|
||||
// TestWaitAgentNonComputerUseNoRecording verifies that when the
|
||||
// child chat is NOT a computer_use chat, no recording is attempted.
|
||||
// StartDesktopRecording must never be called.
|
||||
func TestWaitAgentNonComputerUseNoRecording(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := chatdTestContext(t)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
|
||||
user, model := seedInternalChatDeps(ctx, t, db)
|
||||
|
||||
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
||||
|
||||
// Create parent and regular (non-computer_use) child.
|
||||
parent, child := createParentChildChats(ctx, t, server, user, model)
|
||||
|
||||
// Add an assistant message so the report is extracted.
|
||||
insertAssistantMessage(ctx, t, db, child.ID, model.ID, "Done.")
|
||||
|
||||
// Wait for background processing triggered by CreateChat to
|
||||
// settle before setting up the mock agent connection.
|
||||
server.inflight.Wait()
|
||||
|
||||
// Wire up the mock agent connection. The mock has zero
|
||||
// expectations — gomock will fail if StartDesktopRecording
|
||||
// or any other method is called.
|
||||
server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
||||
return mockConn, func() {}, nil
|
||||
}
|
||||
|
||||
setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
|
||||
|
||||
// Invoke wait_agent via the tool closure — the isComputerUseChat
|
||||
// guard should be false, so no recording calls fire.
|
||||
resp, err := invokeWaitAgentTool(ctx, t, server, db, parent.ID, child.ID, 5)
|
||||
require.NoError(t, err)
|
||||
require.False(t, resp.IsError, "expected successful response, got: %s", resp.Content)
|
||||
|
||||
// Parse the response JSON and verify no recording_file_id.
|
||||
var result map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
||||
_, hasRecording := result["recording_file_id"]
|
||||
assert.False(t, hasRecording, "non-computer_use chat should not produce recording_file_id")
|
||||
}
|
||||
|
||||
// TestWaitAgentRecordingStartFails verifies that when
|
||||
// StartDesktopRecording returns an error, the wait_agent flow still
|
||||
// succeeds and no recording_id is produced. StopDesktopRecording
|
||||
// must NOT be called since the recordingID is cleared on start
|
||||
// failure.
|
||||
func TestWaitAgentRecordingStartFails(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := chatdTestContext(t)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
|
||||
user, model := seedInternalChatDeps(ctx, t, db)
|
||||
workspace, _, agent := seedWorkspaceBinding(t, db, user.ID)
|
||||
|
||||
// Create the server WITHOUT agentConnFn so the background
|
||||
// processing of the parent chat doesn't use the mock.
|
||||
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
||||
|
||||
// Create parent + computer_use child.
|
||||
parent, child := createComputerUseParentChild(
|
||||
ctx, t, server, user, model, workspace, agent,
|
||||
"parent-start-fail", "computer-use-start-fail",
|
||||
)
|
||||
|
||||
// Now wire up the mock agent connection.
|
||||
server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
||||
return mockConn, func() {}, nil
|
||||
}
|
||||
|
||||
insertAssistantMessage(ctx, t, db, child.ID, model.ID, "Opened the browser.")
|
||||
setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
|
||||
|
||||
// StartDesktopRecording fails. StopDesktopRecording must NOT
|
||||
// be called — gomock enforces this: any unexpected call fails
|
||||
// the test.
|
||||
mockConn.EXPECT().
|
||||
StartDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(xerrors.New("ffmpeg not found")).
|
||||
Times(1)
|
||||
|
||||
// Invoke wait_agent via the tool closure.
|
||||
resp, err := invokeWaitAgentTool(ctx, t, server, db, parent.ID, child.ID, 5)
|
||||
require.NoError(t, err)
|
||||
require.False(t, resp.IsError, "recording failure is best-effort, tool should succeed")
|
||||
|
||||
// Parse response JSON and assert no recording_file_id.
|
||||
var result map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
||||
_, hasRecording := result["recording_file_id"]
|
||||
assert.False(t, hasRecording, "no recording_file_id when start fails")
|
||||
}
|
||||
|
||||
// TestWaitAgentRecordingStopFails verifies that when
|
||||
// StopDesktopRecording returns an error, the wait_agent flow still
|
||||
// succeeds but no recording_id is produced.
|
||||
func TestWaitAgentRecordingStopFails(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := chatdTestContext(t)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
|
||||
user, model := seedInternalChatDeps(ctx, t, db)
|
||||
workspace, _, agent := seedWorkspaceBinding(t, db, user.ID)
|
||||
|
||||
// Create the server WITHOUT agentConnFn so the background
|
||||
// processing of the parent chat doesn't use the mock.
|
||||
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
||||
|
||||
// Create parent + computer_use child.
|
||||
parent, child := createComputerUseParentChild(
|
||||
ctx, t, server, user, model, workspace, agent,
|
||||
"parent-stop-fail", "computer-use-stop-fail",
|
||||
)
|
||||
|
||||
// Now wire up the mock agent connection.
|
||||
server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
||||
return mockConn, func() {}, nil
|
||||
}
|
||||
|
||||
insertAssistantMessage(ctx, t, db, child.ID, model.ID, "Checked settings.")
|
||||
setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
|
||||
|
||||
// Start succeeds, stop fails.
|
||||
mockConn.EXPECT().
|
||||
StartDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(nil).
|
||||
Times(1)
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(nil, xerrors.New("disk full")).
|
||||
Times(1)
|
||||
|
||||
// Invoke wait_agent via the tool closure.
|
||||
resp, err := invokeWaitAgentTool(ctx, t, server, db, parent.ID, child.ID, 5)
|
||||
require.NoError(t, err)
|
||||
require.False(t, resp.IsError, "recording failure is best-effort, tool should succeed")
|
||||
|
||||
// Parse response JSON and assert no recording_file_id.
|
||||
var result map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
||||
_, hasRecording := result["recording_file_id"]
|
||||
assert.False(t, hasRecording, "no recording_file_id when stop fails")
|
||||
}
|
||||
|
||||
// TestWaitAgentTimeoutLeavesRecordingRunning verifies that when the
|
||||
// subagent times out, StopDesktopRecording is NOT called. The
|
||||
// recording is left running on the agent so the next wait_agent
|
||||
// call continues it seamlessly.
|
||||
func TestWaitAgentTimeoutLeavesRecordingRunning(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
mClock := quartz.NewMock(t)
|
||||
ctx := chatdTestContext(t)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
|
||||
// Use the mock clock server; don't set agentConnFn yet.
|
||||
server := newInternalTestServerWithClock(t, db, ps, chatprovider.ProviderAPIKeys{}, mClock)
|
||||
|
||||
user, model := seedInternalChatDeps(ctx, t, db)
|
||||
workspace, _, agent := seedWorkspaceBinding(t, db, user.ID)
|
||||
|
||||
// Create parent + computer_use child.
|
||||
_, child := createComputerUseParentChild(
|
||||
ctx, t, server, user, model, workspace, agent,
|
||||
"parent-timeout", "computer-use-timeout",
|
||||
)
|
||||
|
||||
// Set child to running so it never completes.
|
||||
setChatStatus(ctx, t, db, child.ID, database.ChatStatusRunning, "")
|
||||
|
||||
// Now wire up the mock agent connection.
|
||||
server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
||||
return mockConn, func() {}, nil
|
||||
}
|
||||
|
||||
// Start recording succeeds.
|
||||
mockConn.EXPECT().
|
||||
StartDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(nil).
|
||||
Times(1)
|
||||
|
||||
// StopDesktopRecording must NOT be called on timeout.
|
||||
// gomock enforces this: any unexpected call fails the test.
|
||||
|
||||
// Trap the timeout timer to know when the function has entered
|
||||
// its poll loop.
|
||||
timerTrap := mClock.Trap().NewTimer("chatd", "subagent_await")
|
||||
|
||||
type toolResult struct {
|
||||
resp fantasy.ToolResponse
|
||||
err error
|
||||
}
|
||||
resultCh := make(chan toolResult, 1)
|
||||
|
||||
// Re-fetch the parent so LastModelConfigID is populated.
|
||||
parentChat, err := db.GetChatByID(ctx, child.ParentChatID.UUID)
|
||||
require.NoError(t, err)
|
||||
|
||||
tools := server.subagentTools(ctx, func() database.Chat { return parentChat })
|
||||
tool := findToolByName(tools, "wait_agent")
|
||||
require.NotNil(t, tool, "wait_agent tool must be present")
|
||||
|
||||
argsJSON, err := json.Marshal(map[string]any{
|
||||
"chat_id": child.ID.String(),
|
||||
"timeout_seconds": 1,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
go func() {
|
||||
resp, runErr := tool.Run(ctx, fantasy.ToolCall{
|
||||
ID: "test-timeout-call",
|
||||
Name: "wait_agent",
|
||||
Input: string(argsJSON),
|
||||
})
|
||||
resultCh <- toolResult{resp: resp, err: runErr}
|
||||
}()
|
||||
|
||||
// Wait for the timer to be created, then release it.
|
||||
timerTrap.MustWait(ctx).MustRelease(ctx)
|
||||
timerTrap.Close()
|
||||
|
||||
// Advance past the 1s timeout.
|
||||
mClock.Advance(time.Second).MustWait(ctx)
|
||||
|
||||
result := testutil.RequireReceive(ctx, t, resultCh)
|
||||
require.NoError(t, result.err)
|
||||
assert.True(t, result.resp.IsError, "expected error response on timeout")
|
||||
assert.Contains(t, result.resp.Content, "timed out")
|
||||
}
|
||||
|
||||
// TestStopAndStoreRecordingOversized verifies that when the recording
|
||||
// data exceeds MaxRecordingSize, stopAndStoreRecording returns an
|
||||
// empty string and does NOT call InsertChatFile.
|
||||
func TestStopAndStoreRecordingOversized(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := chatdTestContext(t)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
|
||||
user, _ := seedInternalChatDeps(ctx, t, db)
|
||||
workspace, _, _ := seedWorkspaceBinding(t, db, user.ID)
|
||||
|
||||
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
||||
|
||||
// Create a reader that produces MaxRecordingSize+1 bytes without
|
||||
// allocating the full buffer in memory.
|
||||
oversizedReader := io.LimitReader(
|
||||
&zeroReader{},
|
||||
int64(workspacesdk.MaxRecordingSize+1),
|
||||
)
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(io.NopCloser(oversizedReader), nil).
|
||||
Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
storedFileID := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
assert.Empty(t, storedFileID, "oversized recording should not be stored")
|
||||
}
|
||||
|
||||
// TestStopAndStoreRecordingEmpty verifies that when the recording
|
||||
// data is empty, stopAndStoreRecording returns an empty string and
|
||||
// does NOT call InsertChatFile.
|
||||
func TestStopAndStoreRecordingEmpty(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := chatdTestContext(t)
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
|
||||
user, _ := seedInternalChatDeps(ctx, t, db)
|
||||
workspace, _, _ := seedWorkspaceBinding(t, db, user.ID)
|
||||
|
||||
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
||||
|
||||
// Return empty data.
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(io.NopCloser(bytes.NewReader(nil)), nil).
|
||||
Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
storedFileID := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
assert.Empty(t, storedFileID, "empty recording should not be stored")
|
||||
}
|
||||
+80
-10
@@ -4,6 +4,7 @@ import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -12,11 +13,13 @@ import (
|
||||
"github.com/google/uuid"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
)
|
||||
|
||||
var ErrSubagentNotDescendant = xerrors.New("target chat is not a descendant of current chat")
|
||||
@@ -166,22 +169,89 @@ func (p *Server) subagentTools(ctx context.Context, currentChat func() database.
|
||||
}
|
||||
|
||||
parent := currentChat()
|
||||
targetChat, report, err := p.awaitSubagentCompletion(
|
||||
ctx,
|
||||
parent.ID,
|
||||
targetChatID,
|
||||
timeout,
|
||||
)
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(err.Error()), nil
|
||||
|
||||
// Authorize: the target chat must be a descendant
|
||||
// of the current (parent) chat.
|
||||
isDescendant, descErr := isSubagentDescendant(ctx, p.db, parent.ID, targetChatID)
|
||||
if descErr != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
fmt.Sprintf("failed to verify subagent relationship: %v", descErr)), nil
|
||||
}
|
||||
if !isDescendant {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
"target chat is not a subagent of the current chat"), nil
|
||||
}
|
||||
|
||||
return toolJSONResponse(map[string]any{
|
||||
// Check if the target is a computer_use subagent
|
||||
// and start a desktop recording. Failures are
|
||||
// best-effort warnings — recording never blocks
|
||||
// the wait_agent flow.
|
||||
var recordingID string
|
||||
var agentConn workspacesdk.AgentConn
|
||||
|
||||
targetChatInfo, lookupErr := p.db.GetChatByID(ctx, targetChatID)
|
||||
if lookupErr != nil && !xerrors.Is(lookupErr, sql.ErrNoRows) {
|
||||
p.logger.Warn(ctx, "unexpected error looking up chat for recording",
|
||||
slog.F("chat_id", targetChatID),
|
||||
slog.Error(lookupErr),
|
||||
)
|
||||
}
|
||||
isComputerUseChat := lookupErr == nil && targetChatInfo.Mode.Valid &&
|
||||
targetChatInfo.Mode.ChatMode == database.ChatModeComputerUse &&
|
||||
targetChatInfo.AgentID.Valid
|
||||
canRecord := isComputerUseChat && p.agentConnFn != nil
|
||||
|
||||
if canRecord {
|
||||
conn, closeFn, connErr := p.agentConnFn(ctx, targetChatInfo.AgentID.UUID)
|
||||
if connErr == nil {
|
||||
agentConn = conn
|
||||
defer closeFn()
|
||||
|
||||
recordingID = targetChatID.String()
|
||||
startErr := conn.StartDesktopRecording(ctx,
|
||||
workspacesdk.StartDesktopRecordingRequest{RecordingID: recordingID})
|
||||
if startErr != nil {
|
||||
p.logger.Warn(ctx, "failed to start desktop recording",
|
||||
slog.Error(startErr))
|
||||
recordingID = "" // Don't try to stop.
|
||||
}
|
||||
} else {
|
||||
p.logger.Warn(ctx, "failed to get agent conn for recording",
|
||||
slog.Error(connErr))
|
||||
}
|
||||
}
|
||||
|
||||
targetChat, report, awaitErr := p.awaitSubagentCompletion(
|
||||
ctx, parent.ID, targetChatID, timeout,
|
||||
)
|
||||
|
||||
// On timeout/error, leave the recording running on
|
||||
// the agent so the next wait_agent call continues
|
||||
// it seamlessly.
|
||||
if awaitErr != nil {
|
||||
return fantasy.NewTextErrorResponse(awaitErr.Error()), nil
|
||||
}
|
||||
|
||||
// Only stop and store the recording on success.
|
||||
var storedFileID string
|
||||
if recordingID != "" && agentConn != nil {
|
||||
// Use a fresh context for cleanup so a canceled
|
||||
// parent context doesn't prevent recording storage.
|
||||
stopCtx, stopCancel := context.WithTimeout(context.WithoutCancel(ctx), 90*time.Second)
|
||||
defer stopCancel()
|
||||
storedFileID = p.stopAndStoreRecording(stopCtx, agentConn,
|
||||
recordingID, parent.OwnerID, parent.WorkspaceID)
|
||||
}
|
||||
resp := map[string]any{
|
||||
"chat_id": targetChatID.String(),
|
||||
"title": targetChat.Title,
|
||||
"report": report,
|
||||
"status": string(targetChat.Status),
|
||||
}), nil
|
||||
}
|
||||
if storedFileID != "" {
|
||||
resp["recording_file_id"] = storedFileID
|
||||
}
|
||||
return toolJSONResponse(resp), nil
|
||||
},
|
||||
),
|
||||
fantasy.NewAgentTool(
|
||||
|
||||
@@ -94,6 +94,8 @@ type AgentConn interface {
|
||||
WatchGit(ctx context.Context, logger slog.Logger, chatID uuid.UUID) (*wsjson.Stream[codersdk.WorkspaceAgentGitServerMessage, codersdk.WorkspaceAgentGitClientMessage], error)
|
||||
ConnectDesktopVNC(ctx context.Context) (net.Conn, error)
|
||||
ExecuteDesktopAction(ctx context.Context, action DesktopAction) (DesktopActionResponse, error)
|
||||
StartDesktopRecording(ctx context.Context, req StartDesktopRecordingRequest) error
|
||||
StopDesktopRecording(ctx context.Context, req StopDesktopRecordingRequest) (io.ReadCloser, error)
|
||||
}
|
||||
|
||||
// AgentConn represents a connection to a workspace agent.
|
||||
@@ -596,6 +598,23 @@ type DesktopActionResponse struct {
|
||||
ScreenshotHeight int `json:"screenshot_height,omitempty"`
|
||||
}
|
||||
|
||||
// StartDesktopRecordingRequest is the request body for starting a
|
||||
// desktop recording session.
|
||||
type StartDesktopRecordingRequest struct {
|
||||
RecordingID string `json:"recording_id"`
|
||||
}
|
||||
|
||||
// StopDesktopRecordingRequest is the request body for stopping a
|
||||
// desktop recording session.
|
||||
type StopDesktopRecordingRequest struct {
|
||||
RecordingID string `json:"recording_id"`
|
||||
}
|
||||
|
||||
// MaxRecordingSize is the largest desktop recording (in bytes)
|
||||
// that will be accepted. Used by both the agent-side stop handler
|
||||
// and the server-side storage pipeline.
|
||||
const MaxRecordingSize = 100 << 20 // 100 MB
|
||||
|
||||
// ExecuteDesktopAction executes a mouse/keyboard/scroll action on the
|
||||
// agent's desktop.
|
||||
func (c *agentConn) ExecuteDesktopAction(ctx context.Context, action DesktopAction) (DesktopActionResponse, error) {
|
||||
@@ -643,6 +662,43 @@ func (c *agentConn) ExecuteDesktopAction(ctx context.Context, action DesktopActi
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// StartDesktopRecording starts a desktop recording session on the
|
||||
// agent with the given recording ID. The recording ID is
|
||||
// caller-provided and must be unique. Idempotent — if the ID is
|
||||
// already recording, returns success.
|
||||
func (c *agentConn) StartDesktopRecording(ctx context.Context, req StartDesktopRecordingRequest) error {
|
||||
ctx, span := tracing.StartSpan(ctx)
|
||||
defer span.End()
|
||||
res, err := c.apiRequest(ctx, http.MethodPost, "/api/v0/desktop/recording/start", req)
|
||||
if err != nil {
|
||||
return xerrors.Errorf("start recording request: %w", err)
|
||||
}
|
||||
defer res.Body.Close()
|
||||
if res.StatusCode != http.StatusOK {
|
||||
return codersdk.ReadBodyAsError(res)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// StopDesktopRecording stops a desktop recording session on the
|
||||
// agent and returns the MP4 data as an io.ReadCloser. The caller
|
||||
// is responsible for closing the returned reader. Idempotent —
|
||||
// safe to call on an already-stopped recording.
|
||||
func (c *agentConn) StopDesktopRecording(ctx context.Context, req StopDesktopRecordingRequest) (io.ReadCloser, error) {
|
||||
ctx, span := tracing.StartSpan(ctx)
|
||||
defer span.End()
|
||||
res, err := c.apiRequest(ctx, http.MethodPost, "/api/v0/desktop/recording/stop", req)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("stop recording request: %w", err)
|
||||
}
|
||||
if res.StatusCode != http.StatusOK {
|
||||
defer res.Body.Close()
|
||||
return nil, codersdk.ReadBodyAsError(res)
|
||||
}
|
||||
// Caller is responsible for closing res.Body.
|
||||
return res.Body, nil
|
||||
}
|
||||
|
||||
// DeleteDevcontainer deletes the provided devcontainer.
|
||||
// This is a blocking call and will wait for the container to be deleted.
|
||||
func (c *agentConn) DeleteDevcontainer(ctx context.Context, devcontainerID string) error {
|
||||
|
||||
@@ -550,6 +550,20 @@ func (mr *MockAgentConnMockRecorder) Speedtest(ctx, direction, duration any) *go
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Speedtest", reflect.TypeOf((*MockAgentConn)(nil).Speedtest), ctx, direction, duration)
|
||||
}
|
||||
|
||||
// StartDesktopRecording mocks base method.
|
||||
func (m *MockAgentConn) StartDesktopRecording(ctx context.Context, req workspacesdk.StartDesktopRecordingRequest) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "StartDesktopRecording", ctx, req)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// StartDesktopRecording indicates an expected call of StartDesktopRecording.
|
||||
func (mr *MockAgentConnMockRecorder) StartDesktopRecording(ctx, req any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StartDesktopRecording", reflect.TypeOf((*MockAgentConn)(nil).StartDesktopRecording), ctx, req)
|
||||
}
|
||||
|
||||
// StartProcess mocks base method.
|
||||
func (m *MockAgentConn) StartProcess(ctx context.Context, req workspacesdk.StartProcessRequest) (workspacesdk.StartProcessResponse, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -565,6 +579,21 @@ func (mr *MockAgentConnMockRecorder) StartProcess(ctx, req any) *gomock.Call {
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StartProcess", reflect.TypeOf((*MockAgentConn)(nil).StartProcess), ctx, req)
|
||||
}
|
||||
|
||||
// StopDesktopRecording mocks base method.
|
||||
func (m *MockAgentConn) StopDesktopRecording(ctx context.Context, req workspacesdk.StopDesktopRecordingRequest) (io.ReadCloser, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "StopDesktopRecording", ctx, req)
|
||||
ret0, _ := ret[0].(io.ReadCloser)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// StopDesktopRecording indicates an expected call of StopDesktopRecording.
|
||||
func (mr *MockAgentConnMockRecorder) StopDesktopRecording(ctx, req any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StopDesktopRecording", reflect.TypeOf((*MockAgentConn)(nil).StopDesktopRecording), ctx, req)
|
||||
}
|
||||
|
||||
// TailnetConn mocks base method.
|
||||
func (m *MockAgentConn) TailnetConn() *tailnet.Conn {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
Reference in New Issue
Block a user