mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: agents desktop recording thumbnail backend (#24022)
The agents chat interface displays thumbnails for videos recorded by the computer use agent. Currently, to display a thumbnail, the frontend downloads the entire video and shows the first frame. This PR starts storing a new thumbnail file in the database for every recorded video, and exposes the file id in the `wait_agent` tool result alongside the recording file id, so the frontend can fetch just the thumbnail.
This commit is contained in:
Hugo Dutka
@@ -73,9 +73,9 @@ const (
|
||||
|
||||
// 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.
|
||||
// slot buffers up to MaxRecordingSize + MaxThumbnailSize
|
||||
// (110 MB) in memory, so this value implicitly bounds memory
|
||||
// to roughly maxConcurrentRecordingUploads * 110 MB.
|
||||
maxConcurrentRecordingUploads = 25
|
||||
|
||||
// staleRecoveryIntervalDivisor determines how often the stale
|
||||
|
||||
+148
-54
@@ -2,8 +2,11 @@ package chatd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/multipart"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
@@ -13,71 +16,60 @@ import (
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
)
|
||||
|
||||
type recordingResult struct {
|
||||
recordingFileID string
|
||||
thumbnailFileID string
|
||||
}
|
||||
|
||||
// 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.
|
||||
// multipart response containing the MP4 and optional thumbnail, and
|
||||
// stores them in chat_files. Only called when the subagent completed
|
||||
// successfully. Returns file IDs on success, empty fields 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 {
|
||||
) recordingResult {
|
||||
var result recordingResult
|
||||
|
||||
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 ""
|
||||
return result
|
||||
}
|
||||
|
||||
body, err := conn.StopDesktopRecording(ctx,
|
||||
resp, err := conn.StopDesktopRecording(ctx,
|
||||
workspacesdk.StopDesktopRecordingRequest{RecordingID: recordingID})
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "failed to stop desktop recording",
|
||||
slog.Error(err))
|
||||
return ""
|
||||
return result
|
||||
}
|
||||
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}
|
||||
}()
|
||||
defer resp.Body.Close()
|
||||
|
||||
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 ""
|
||||
_, params, err := mime.ParseMediaType(resp.ContentType)
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "failed to parse content type from recording response",
|
||||
slog.F("content_type", resp.ContentType),
|
||||
slog.Error(err))
|
||||
return result
|
||||
}
|
||||
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 ""
|
||||
boundary := params["boundary"]
|
||||
if boundary == "" {
|
||||
p.logger.Warn(ctx, "missing boundary in recording response content type",
|
||||
slog.F("content_type", resp.ContentType))
|
||||
return result
|
||||
}
|
||||
|
||||
if !workspaceID.Valid {
|
||||
p.logger.Warn(ctx, "chat has no workspace, cannot store recording")
|
||||
return ""
|
||||
return result
|
||||
}
|
||||
|
||||
// The chatd actor is used here because the recording is stored on
|
||||
@@ -87,21 +79,123 @@ func (p *Server) stopAndStoreRecording(
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "failed to resolve workspace for recording",
|
||||
slog.Error(err))
|
||||
return ""
|
||||
return result
|
||||
}
|
||||
|
||||
//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 ""
|
||||
mr := multipart.NewReader(resp.Body, boundary)
|
||||
// Context cancellation is checked between parts. Within a
|
||||
// part read, cancellation relies on Go's HTTP transport closing
|
||||
// the underlying connection when the context is done, which
|
||||
// interrupts the blocked io.ReadAll.
|
||||
// First pass: parse all multipart parts into memory.
|
||||
// The agent sends at most two parts: one video/mp4 and one
|
||||
// optional image/jpeg thumbnail. Cap the number of parts to
|
||||
// prevent a malicious or broken agent from forcing the server
|
||||
// into an unbounded parsing loop.
|
||||
const maxParts = 2
|
||||
var videoData, thumbnailData []byte
|
||||
for range maxParts {
|
||||
if ctx.Err() != nil {
|
||||
p.logger.Warn(ctx, "context canceled while reading recording parts", slog.Error(ctx.Err()))
|
||||
break
|
||||
}
|
||||
|
||||
part, err := mr.NextPart()
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "error reading next multipart part", slog.Error(err))
|
||||
break
|
||||
}
|
||||
|
||||
contentType := part.Header.Get("Content-Type")
|
||||
|
||||
// Select the read limit based on content type so that
|
||||
// thumbnails (image/jpeg) do not allocate up to
|
||||
// MaxRecordingSize (100 MB) before the size check rejects
|
||||
// them. Unknown types use a small default since they are
|
||||
// discarded below.
|
||||
maxSize := int64(1 << 20) // 1 MB default for unknown types
|
||||
switch contentType {
|
||||
case "video/mp4":
|
||||
maxSize = int64(workspacesdk.MaxRecordingSize)
|
||||
case "image/jpeg":
|
||||
maxSize = int64(workspacesdk.MaxThumbnailSize)
|
||||
}
|
||||
|
||||
data, err := io.ReadAll(io.LimitReader(part, maxSize+1))
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "failed to read recording part data",
|
||||
slog.F("content_type", contentType),
|
||||
slog.Error(err))
|
||||
continue
|
||||
}
|
||||
if int64(len(data)) > maxSize {
|
||||
p.logger.Warn(ctx, "recording part exceeds maximum size, skipping",
|
||||
slog.F("content_type", contentType),
|
||||
slog.F("size", len(data)),
|
||||
slog.F("max_size", maxSize))
|
||||
continue
|
||||
}
|
||||
if len(data) == 0 {
|
||||
p.logger.Warn(ctx, "recording part is empty, skipping",
|
||||
slog.F("content_type", contentType))
|
||||
continue
|
||||
}
|
||||
|
||||
switch contentType {
|
||||
case "video/mp4":
|
||||
if videoData != nil {
|
||||
p.logger.Warn(ctx, "duplicate video/mp4 part in recording response, skipping")
|
||||
continue
|
||||
}
|
||||
videoData = data
|
||||
case "image/jpeg":
|
||||
if thumbnailData != nil {
|
||||
p.logger.Warn(ctx, "duplicate image/jpeg part in recording response, skipping")
|
||||
continue
|
||||
}
|
||||
thumbnailData = data
|
||||
default:
|
||||
p.logger.Debug(ctx, "skipping unknown part content type",
|
||||
slog.F("content_type", contentType))
|
||||
}
|
||||
}
|
||||
return row.ID.String()
|
||||
|
||||
// Second pass: store the collected data in the database.
|
||||
if videoData != nil {
|
||||
//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: videoData,
|
||||
})
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "failed to store recording in database",
|
||||
slog.Error(err))
|
||||
} else {
|
||||
result.recordingFileID = row.ID.String()
|
||||
}
|
||||
}
|
||||
if thumbnailData != nil && result.recordingFileID != "" {
|
||||
//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("thumbnail-%s.jpg", p.clock.Now().UTC().Format("2006-01-02T15-04-05Z")),
|
||||
Mimetype: "image/jpeg",
|
||||
Data: thumbnailData,
|
||||
})
|
||||
if err != nil {
|
||||
p.logger.Warn(ctx, "failed to store thumbnail in database",
|
||||
slog.Error(err))
|
||||
} else {
|
||||
result.thumbnailFileID = row.ID.String()
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
@@ -5,6 +5,8 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/textproto"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -34,6 +36,30 @@ func (zeroReader) Read(p []byte) (int, error) {
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
// partSpec describes a single part for buildMultipartResponse.
|
||||
type partSpec struct {
|
||||
contentType string
|
||||
data []byte
|
||||
}
|
||||
|
||||
// buildMultipartResponse constructs a StopDesktopRecordingResponse
|
||||
// with the given content type/data pairs encoded as multipart/mixed.
|
||||
func buildMultipartResponse(parts ...partSpec) workspacesdk.StopDesktopRecordingResponse {
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
for _, p := range parts {
|
||||
partWriter, _ := mw.CreatePart(textproto.MIMEHeader{
|
||||
"Content-Type": {p.contentType},
|
||||
})
|
||||
_, _ = partWriter.Write(p.data)
|
||||
}
|
||||
_ = mw.Close()
|
||||
return workspacesdk.StopDesktopRecordingResponse{
|
||||
Body: io.NopCloser(bytes.NewReader(buf.Bytes())),
|
||||
ContentType: "multipart/mixed; boundary=" + mw.Boundary(),
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
@@ -170,8 +196,7 @@ func TestWaitAgentComputerUseRecording(t *testing.T) {
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(io.NopCloser(bytes.NewReader(fakeMp4)), nil).
|
||||
Times(1)
|
||||
Return(buildMultipartResponse(partSpec{"video/mp4", fakeMp4}), nil).Times(1)
|
||||
|
||||
// Invoke wait_agent via the tool closure.
|
||||
resp, err := invokeWaitAgentTool(ctx, t, server, db, parent.ID, child.ID, 5)
|
||||
@@ -198,6 +223,87 @@ func TestWaitAgentComputerUseRecording(t *testing.T) {
|
||||
assert.Equal(t, fakeMp4, chatFile.Data)
|
||||
}
|
||||
|
||||
// TestWaitAgentComputerUseRecordingWithThumbnail verifies the
|
||||
// recording flow when the agent produces both video and thumbnail:
|
||||
// both file IDs appear in the wait_agent tool response.
|
||||
func TestWaitAgentComputerUseRecordingWithThumbnail(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)
|
||||
|
||||
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
||||
|
||||
parent, child := createComputerUseParentChild(
|
||||
ctx, t, server, user, model, workspace, agent,
|
||||
"parent-recording-thumb", "computer-use-child-thumb",
|
||||
)
|
||||
|
||||
server.drainInflight()
|
||||
|
||||
server.agentConnFn = func(_ context.Context, agentID uuid.UUID) (workspacesdk.AgentConn, func(), error) {
|
||||
require.Equal(t, agent.ID, agentID)
|
||||
return mockConn, func() {}, nil
|
||||
}
|
||||
|
||||
insertAssistantMessage(ctx, t, db, child.ID, model.ID, "I opened Firefox and took a screenshot.")
|
||||
|
||||
setChatStatus(ctx, t, db, child.ID, database.ChatStatusWaiting, "")
|
||||
|
||||
fakeMp4 := []byte("fake-mp4-data-with-thumbnail-test")
|
||||
fakeThumb := []byte("fake-jpeg-thumbnail-data")
|
||||
|
||||
mockConn.EXPECT().
|
||||
StartDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
DoAndReturn(func(_ context.Context, req workspacesdk.StartDesktopRecordingRequest) error {
|
||||
require.NotEmpty(t, req.RecordingID)
|
||||
return nil
|
||||
}).
|
||||
Times(1)
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(buildMultipartResponse(
|
||||
partSpec{"video/mp4", fakeMp4},
|
||||
partSpec{"image/jpeg", fakeThumb},
|
||||
), nil).Times(1)
|
||||
|
||||
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)
|
||||
|
||||
var result map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
||||
|
||||
// Verify recording_file_id is present and valid.
|
||||
storedFileID, ok := result["recording_file_id"].(string)
|
||||
require.True(t, ok, "recording_file_id must be present in response")
|
||||
require.NotEmpty(t, storedFileID)
|
||||
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.Equal(t, fakeMp4, chatFile.Data)
|
||||
|
||||
// Verify thumbnail_file_id is present and valid.
|
||||
thumbFileID, ok := result["thumbnail_file_id"].(string)
|
||||
require.True(t, ok, "thumbnail_file_id must be present in response")
|
||||
require.NotEmpty(t, thumbFileID)
|
||||
thumbUUID, err := uuid.Parse(thumbFileID)
|
||||
require.NoError(t, err)
|
||||
thumbFile, err := db.GetChatFileByID(ctx, thumbUUID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "image/jpeg", thumbFile.Mimetype)
|
||||
assert.Equal(t, fakeThumb, thumbFile.Data)
|
||||
}
|
||||
|
||||
// TestWaitAgentNonComputerUseNoRecording verifies that when the
|
||||
// child chat is NOT a computer_use chat, no recording is attempted.
|
||||
// StartDesktopRecording must never be called.
|
||||
@@ -342,7 +448,7 @@ func TestWaitAgentRecordingStopFails(t *testing.T) {
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(nil, xerrors.New("disk full")).
|
||||
Return(workspacesdk.StopDesktopRecordingResponse{}, xerrors.New("disk full")).
|
||||
Times(1)
|
||||
|
||||
// Invoke wait_agent via the tool closure.
|
||||
@@ -446,10 +552,10 @@ func TestWaitAgentTimeoutLeavesRecordingRunning(t *testing.T) {
|
||||
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) {
|
||||
// TestStopAndStoreRecording_Oversized verifies that when the
|
||||
// recording data exceeds MaxRecordingSize, stopAndStoreRecording
|
||||
// returns an empty string and does NOT call InsertChatFile.
|
||||
func TestStopAndStoreRecording_Oversized(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
@@ -463,29 +569,146 @@ func TestStopAndStoreRecordingOversized(t *testing.T) {
|
||||
|
||||
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),
|
||||
)
|
||||
// Build a streaming multipart response with a video/mp4 part
|
||||
// that exceeds MaxRecordingSize without allocating the full
|
||||
// buffer in memory.
|
||||
pr, pw := io.Pipe()
|
||||
mw := multipart.NewWriter(pw)
|
||||
go func() {
|
||||
partWriter, _ := mw.CreatePart(textproto.MIMEHeader{
|
||||
"Content-Type": {"video/mp4"},
|
||||
})
|
||||
// Stream MaxRecordingSize+1 zero bytes.
|
||||
_, _ = io.Copy(partWriter, io.LimitReader(&zeroReader{}, int64(workspacesdk.MaxRecordingSize+1)))
|
||||
_ = mw.Close()
|
||||
_ = pw.Close()
|
||||
}()
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(io.NopCloser(oversizedReader), nil).
|
||||
Return(workspacesdk.StopDesktopRecordingResponse{
|
||||
Body: pr,
|
||||
ContentType: "multipart/mixed; boundary=" + mw.Boundary(),
|
||||
}, nil).
|
||||
Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
storedFileID := server.stopAndStoreRecording(
|
||||
result := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
assert.Empty(t, storedFileID, "oversized recording should not be stored")
|
||||
assert.Empty(t, result.recordingFileID, "oversized recording should not be stored")
|
||||
}
|
||||
|
||||
// TestStopAndStoreRecordingEmpty verifies that when the recording
|
||||
// TestStopAndStoreRecording_OversizedThumbnail verifies that when the
|
||||
// thumbnail part exceeds MaxThumbnailSize it is skipped while the
|
||||
// normal-sized video part is still stored.
|
||||
func TestStopAndStoreRecording_OversizedThumbnail(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{})
|
||||
|
||||
videoData := bytes.Repeat([]byte{0xAA}, 1024)
|
||||
|
||||
// Build a streaming multipart response with a normal video part
|
||||
// and an oversized thumbnail part.
|
||||
pr, pw := io.Pipe()
|
||||
mw := multipart.NewWriter(pw)
|
||||
go func() {
|
||||
vw, _ := mw.CreatePart(textproto.MIMEHeader{
|
||||
"Content-Type": {"video/mp4"},
|
||||
})
|
||||
_, _ = vw.Write(videoData)
|
||||
tw, _ := mw.CreatePart(textproto.MIMEHeader{
|
||||
"Content-Type": {"image/jpeg"},
|
||||
})
|
||||
// Stream MaxThumbnailSize+1 zero bytes for the thumbnail.
|
||||
_, _ = io.Copy(tw, io.LimitReader(&zeroReader{}, int64(workspacesdk.MaxThumbnailSize+1)))
|
||||
_ = mw.Close()
|
||||
_ = pw.Close()
|
||||
}()
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(workspacesdk.StopDesktopRecordingResponse{
|
||||
Body: pr,
|
||||
ContentType: "multipart/mixed; boundary=" + mw.Boundary(),
|
||||
}, nil).
|
||||
Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
result := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
|
||||
// Video should be stored.
|
||||
recUUID, err := uuid.Parse(result.recordingFileID)
|
||||
require.NoError(t, err, "RecordingFileID should be a valid UUID")
|
||||
recFile, err := db.GetChatFileByID(ctx, recUUID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "video/mp4", recFile.Mimetype)
|
||||
assert.Equal(t, videoData, recFile.Data)
|
||||
|
||||
// Thumbnail should be skipped (oversized).
|
||||
assert.Empty(t, result.thumbnailFileID, "oversized thumbnail should not be stored")
|
||||
}
|
||||
|
||||
// TestStopAndStoreRecording_DuplicatePartsIgnored verifies that when
|
||||
// a multipart response contains two video/mp4 parts, only the first
|
||||
// is stored and the duplicate is skipped.
|
||||
func TestStopAndStoreRecording_DuplicatePartsIgnored(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{})
|
||||
|
||||
firstVideo := bytes.Repeat([]byte{0x01}, 512)
|
||||
secondVideo := bytes.Repeat([]byte{0x02}, 512)
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(buildMultipartResponse(
|
||||
partSpec{"video/mp4", firstVideo},
|
||||
partSpec{"video/mp4", secondVideo},
|
||||
), nil).
|
||||
Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
result := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
|
||||
// Only the first video part should be stored.
|
||||
recUUID, err := uuid.Parse(result.recordingFileID)
|
||||
require.NoError(t, err)
|
||||
recFile, err := db.GetChatFileByID(ctx, recUUID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, firstVideo, recFile.Data, "first video part should be stored, not the duplicate")
|
||||
}
|
||||
|
||||
// TestStopAndStoreRecording_Empty verifies that when the recording
|
||||
// data is empty, stopAndStoreRecording returns an empty string and
|
||||
// does NOT call InsertChatFile.
|
||||
func TestStopAndStoreRecordingEmpty(t *testing.T) {
|
||||
func TestStopAndStoreRecording_Empty(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
@@ -499,16 +722,265 @@ func TestStopAndStoreRecordingEmpty(t *testing.T) {
|
||||
|
||||
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
||||
|
||||
// Return empty data.
|
||||
// Build a multipart response with an empty video/mp4 part.
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(io.NopCloser(bytes.NewReader(nil)), nil).
|
||||
Times(1)
|
||||
Return(buildMultipartResponse(partSpec{"video/mp4", nil}), nil).Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
storedFileID := server.stopAndStoreRecording(
|
||||
result := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
assert.Empty(t, storedFileID, "empty recording should not be stored")
|
||||
assert.Empty(t, result.recordingFileID, "empty recording should not be stored")
|
||||
}
|
||||
|
||||
// TestStopAndStoreRecording_WithThumbnail verifies that a multipart
|
||||
// response containing both a video/mp4 part and an image/jpeg part
|
||||
// results in both files being stored with correct mimetypes.
|
||||
func TestStopAndStoreRecording_WithThumbnail(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{})
|
||||
|
||||
videoData := bytes.Repeat([]byte{0xDE, 0xAD}, 512) // 1024 bytes
|
||||
thumbData := bytes.Repeat([]byte{0xFF, 0xD8}, 256) // 512 bytes
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(buildMultipartResponse(
|
||||
partSpec{"video/mp4", videoData},
|
||||
partSpec{"image/jpeg", thumbData},
|
||||
), nil).
|
||||
Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
result := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
|
||||
// Both file IDs should be valid UUIDs.
|
||||
recUUID, err := uuid.Parse(result.recordingFileID)
|
||||
require.NoError(t, err, "RecordingFileID should be a valid UUID")
|
||||
|
||||
thumbUUID, err := uuid.Parse(result.thumbnailFileID)
|
||||
require.NoError(t, err, "ThumbnailFileID should be a valid UUID")
|
||||
// Verify the recording file in the database.
|
||||
recFile, err := db.GetChatFileByID(ctx, recUUID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "video/mp4", recFile.Mimetype)
|
||||
assert.Equal(t, videoData, recFile.Data)
|
||||
|
||||
// Verify the thumbnail file in the database.
|
||||
thumbFile, err := db.GetChatFileByID(ctx, thumbUUID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "image/jpeg", thumbFile.Mimetype)
|
||||
assert.Equal(t, thumbData, thumbFile.Data)
|
||||
}
|
||||
|
||||
// TestStopAndStoreRecording_VideoOnly verifies that a multipart
|
||||
// response with only a video/mp4 part stores the recording but
|
||||
// leaves thumbnailFileID empty.
|
||||
func TestStopAndStoreRecording_VideoOnly(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{})
|
||||
|
||||
videoData := make([]byte, 1024)
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(buildMultipartResponse(partSpec{"video/mp4", videoData}), nil).Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
result := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
|
||||
// Recording should be stored.
|
||||
recUUID, err := uuid.Parse(result.recordingFileID)
|
||||
require.NoError(t, err, "RecordingFileID should be a valid UUID")
|
||||
|
||||
recFile, err := db.GetChatFileByID(ctx, recUUID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "video/mp4", recFile.Mimetype)
|
||||
assert.Equal(t, videoData, recFile.Data)
|
||||
|
||||
// No thumbnail.
|
||||
assert.Empty(t, result.thumbnailFileID, "ThumbnailFileID should be empty when no thumbnail part is present")
|
||||
}
|
||||
|
||||
// TestStopAndStoreRecording_DownloadFailure verifies that when
|
||||
// StopDesktopRecording returns an error, stopAndStoreRecording
|
||||
// returns an empty recordingResult without panicking.
|
||||
func TestStopAndStoreRecording_DownloadFailure(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{})
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(workspacesdk.StopDesktopRecordingResponse{}, xerrors.New("network error")).
|
||||
Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
result := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
|
||||
assert.Empty(t, result.recordingFileID, "RecordingFileID should be empty on download failure")
|
||||
assert.Empty(t, result.thumbnailFileID, "ThumbnailFileID should be empty on download failure")
|
||||
}
|
||||
|
||||
// TestStopAndStoreRecording_UnknownPartIgnored verifies that parts
|
||||
// with unrecognized content types are silently skipped while known
|
||||
// parts (video/mp4 and image/jpeg) are still stored.
|
||||
func TestStopAndStoreRecording_UnknownPartIgnored(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{})
|
||||
|
||||
videoData := make([]byte, 1024)
|
||||
thumbData := make([]byte, 512)
|
||||
unknownData := make([]byte, 256)
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(buildMultipartResponse(
|
||||
partSpec{"video/mp4", videoData},
|
||||
partSpec{"image/jpeg", thumbData},
|
||||
partSpec{"application/octet-stream", unknownData},
|
||||
), nil).Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
result := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
|
||||
// Both known parts should be stored.
|
||||
recUUID, err := uuid.Parse(result.recordingFileID)
|
||||
require.NoError(t, err, "RecordingFileID should be a valid UUID")
|
||||
|
||||
thumbUUID, err := uuid.Parse(result.thumbnailFileID)
|
||||
require.NoError(t, err, "ThumbnailFileID should be a valid UUID")
|
||||
|
||||
// Verify only 2 files exist (unknown part was skipped).
|
||||
recFile, err := db.GetChatFileByID(ctx, recUUID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "video/mp4", recFile.Mimetype)
|
||||
assert.Equal(t, videoData, recFile.Data)
|
||||
|
||||
thumbFile, err := db.GetChatFileByID(ctx, thumbUUID)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "image/jpeg", thumbFile.Mimetype)
|
||||
assert.Equal(t, thumbData, thumbFile.Data)
|
||||
}
|
||||
|
||||
// TestStopAndStoreRecording_MalformedContentType verifies that a
|
||||
// response with an unparseable Content-Type returns an empty result.
|
||||
func TestStopAndStoreRecording_MalformedContentType(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{})
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(workspacesdk.StopDesktopRecordingResponse{
|
||||
Body: io.NopCloser(bytes.NewReader(nil)),
|
||||
ContentType: "",
|
||||
}, nil).
|
||||
Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
result := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
|
||||
assert.Empty(t, result.recordingFileID, "RecordingFileID should be empty for malformed content type")
|
||||
assert.Empty(t, result.thumbnailFileID, "ThumbnailFileID should be empty for malformed content type")
|
||||
}
|
||||
|
||||
// TestStopAndStoreRecording_MissingBoundary verifies that a
|
||||
// multipart response without a boundary parameter returns an empty
|
||||
// result.
|
||||
func TestStopAndStoreRecording_MissingBoundary(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{})
|
||||
|
||||
mockConn.EXPECT().
|
||||
StopDesktopRecording(gomock.Any(), gomock.Any()).
|
||||
Return(workspacesdk.StopDesktopRecordingResponse{
|
||||
Body: io.NopCloser(bytes.NewReader(nil)),
|
||||
ContentType: "multipart/mixed",
|
||||
}, nil).
|
||||
Times(1)
|
||||
|
||||
recordingID := uuid.New().String()
|
||||
result := server.stopAndStoreRecording(
|
||||
ctx, mockConn, recordingID, user.ID,
|
||||
uuid.NullUUID{UUID: workspace.ID, Valid: true},
|
||||
)
|
||||
|
||||
assert.Empty(t, result.recordingFileID, "RecordingFileID should be empty when boundary is missing")
|
||||
assert.Empty(t, result.thumbnailFileID, "ThumbnailFileID should be empty when boundary is missing")
|
||||
}
|
||||
|
||||
@@ -233,13 +233,13 @@ func (p *Server) subagentTools(ctx context.Context, currentChat func() database.
|
||||
}
|
||||
|
||||
// Only stop and store the recording on success.
|
||||
var storedFileID string
|
||||
var recResult recordingResult
|
||||
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,
|
||||
recResult = p.stopAndStoreRecording(stopCtx, agentConn,
|
||||
recordingID, parent.OwnerID, parent.WorkspaceID)
|
||||
}
|
||||
resp := map[string]any{
|
||||
@@ -248,8 +248,11 @@ func (p *Server) subagentTools(ctx context.Context, currentChat func() database.
|
||||
"report": report,
|
||||
"status": string(targetChat.Status),
|
||||
}
|
||||
if storedFileID != "" {
|
||||
resp["recording_file_id"] = storedFileID
|
||||
if recResult.recordingFileID != "" {
|
||||
resp["recording_file_id"] = recResult.recordingFileID
|
||||
}
|
||||
if recResult.thumbnailFileID != "" {
|
||||
resp["thumbnail_file_id"] = recResult.thumbnailFileID
|
||||
}
|
||||
return toolJSONResponse(resp), nil
|
||||
},
|
||||
|
||||
Reference in New Issue
Block a user