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:
Hugo Dutka
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user