mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: chat desktop backend (#23005)
Implement the backend for the desktop feature for agents. - Adds a new `/api/experimental/chats/$id/desktop` endpoint to coderd which exposes a VNC stream from a [portabledesktop](https://github.com/coder/portabledesktop) process running inside the workspace - Adds a new `spawn_computer_use_agent` tool to chatd, which spawns a subagent that has access to the `computer` tool which lets it interact with the `portabledesktop` process running inside the workspace - Adds the plumbing to make the above possible There's a follow up frontend PR here: https://github.com/coder/coder/pull/23006
This commit is contained in:
Generated
+43
@@ -481,6 +481,49 @@ const docTemplate = `{
|
||||
}
|
||||
}
|
||||
},
|
||||
"/chats/{chat}/archive": {
|
||||
"post": {
|
||||
"tags": [
|
||||
"Chats"
|
||||
],
|
||||
"summary": "Archive a chat",
|
||||
"operationId": "archive-chat",
|
||||
"responses": {
|
||||
"204": {
|
||||
"description": "No Content"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/chats/{chat}/desktop": {
|
||||
"get": {
|
||||
"security": [
|
||||
{
|
||||
"CoderSessionToken": []
|
||||
}
|
||||
],
|
||||
"tags": [
|
||||
"Chats"
|
||||
],
|
||||
"summary": "Watch chat desktop",
|
||||
"operationId": "watch-chat-desktop",
|
||||
"parameters": [
|
||||
{
|
||||
"type": "string",
|
||||
"format": "uuid",
|
||||
"description": "Chat ID",
|
||||
"name": "chat",
|
||||
"in": "path",
|
||||
"required": true
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"101": {
|
||||
"description": "Switching Protocols"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/connectionlog": {
|
||||
"get": {
|
||||
"security": [
|
||||
|
||||
Generated
+39
@@ -410,6 +410,45 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"/chats/{chat}/archive": {
|
||||
"post": {
|
||||
"tags": ["Chats"],
|
||||
"summary": "Archive a chat",
|
||||
"operationId": "archive-chat",
|
||||
"responses": {
|
||||
"204": {
|
||||
"description": "No Content"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/chats/{chat}/desktop": {
|
||||
"get": {
|
||||
"security": [
|
||||
{
|
||||
"CoderSessionToken": []
|
||||
}
|
||||
],
|
||||
"tags": ["Chats"],
|
||||
"summary": "Watch chat desktop",
|
||||
"operationId": "watch-chat-desktop",
|
||||
"parameters": [
|
||||
{
|
||||
"type": "string",
|
||||
"format": "uuid",
|
||||
"description": "Chat ID",
|
||||
"name": "chat",
|
||||
"in": "path",
|
||||
"required": true
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"101": {
|
||||
"description": "Switching Protocols"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/connectionlog": {
|
||||
"get": {
|
||||
"security": [
|
||||
|
||||
+58
-18
@@ -32,6 +32,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/webpush"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -179,6 +180,7 @@ type CreateOptions struct {
|
||||
RootChatID uuid.NullUUID
|
||||
Title string
|
||||
ModelConfigID uuid.UUID
|
||||
ChatMode database.NullChatMode
|
||||
SystemPrompt string
|
||||
InitialUserContent []codersdk.ChatMessagePart
|
||||
}
|
||||
@@ -262,6 +264,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C
|
||||
RootChatID: opts.RootChatID,
|
||||
LastModelConfigID: opts.ModelConfigID,
|
||||
Title: opts.Title,
|
||||
Mode: opts.ChatMode,
|
||||
})
|
||||
if err != nil {
|
||||
return xerrors.Errorf("insert chat: %w", err)
|
||||
@@ -2143,10 +2146,13 @@ func (p *Server) runChat(
|
||||
// Fire title generation asynchronously so it doesn't block the
|
||||
// chat response. It uses a detached context so it can finish
|
||||
// even after the chat processing context is canceled.
|
||||
// Snapshot model so the goroutine doesn't race with the
|
||||
// model = cuModel reassignment below.
|
||||
titleModel := model
|
||||
p.inflight.Add(1)
|
||||
go func() {
|
||||
defer p.inflight.Done()
|
||||
p.maybeGenerateChatTitle(context.WithoutCancel(ctx), chat, messages, model, providerKeys, logger)
|
||||
p.maybeGenerateChatTitle(context.WithoutCancel(ctx), chat, messages, titleModel, providerKeys, logger)
|
||||
}()
|
||||
|
||||
prompt, err := chatprompt.ConvertMessagesWithFiles(ctx, messages, p.chatFileResolver(), logger)
|
||||
@@ -2157,6 +2163,9 @@ func (p *Server) runChat(
|
||||
prompt = chatprompt.InsertSystem(prompt, defaultSubagentInstruction)
|
||||
}
|
||||
|
||||
// Detect computer-use subagent via the mode column.
|
||||
isComputerUse := chat.Mode.Valid && chat.Mode.ChatMode == database.ChatModeComputerUse
|
||||
|
||||
// NOTE: Buffering was already started in processChat before
|
||||
// the running status was published, so message_part events
|
||||
// are captured from the moment subscribers can see
|
||||
@@ -2512,6 +2521,20 @@ func (p *Server) runChat(
|
||||
},
|
||||
}
|
||||
|
||||
if isComputerUse {
|
||||
// Override model for computer use subagent.
|
||||
cuModel, cuErr := chatprovider.ModelFromConfig(
|
||||
chattool.ComputerUseModelProvider,
|
||||
chattool.ComputerUseModelName,
|
||||
providerKeys,
|
||||
chatprovider.UserAgent(),
|
||||
)
|
||||
if cuErr != nil {
|
||||
return xerrors.Errorf("resolve computer use model: %w", cuErr)
|
||||
}
|
||||
model = cuModel
|
||||
}
|
||||
|
||||
// Here are all the tools we have for the chat.
|
||||
tools := []fantasy.AgentTool{
|
||||
chattool.ReadFile(chattool.ReadFileOptions{
|
||||
@@ -2568,23 +2591,34 @@ func (p *Server) runChat(
|
||||
WorkspaceMu: &workspaceMu,
|
||||
}),
|
||||
)
|
||||
tools = append(tools, p.subagentTools(func() database.Chat {
|
||||
tools = append(tools, p.subagentTools(ctx, func() database.Chat {
|
||||
return chat
|
||||
})...)
|
||||
}
|
||||
|
||||
// Build provider-native tools (e.g., web search) based on
|
||||
// the model configuration.
|
||||
var providerTools []fantasy.Tool
|
||||
var providerTools []chatloop.ProviderTool
|
||||
if callConfig.ProviderOptions != nil {
|
||||
providerTools = buildProviderTools(model.Provider(), callConfig.ProviderOptions)
|
||||
}
|
||||
|
||||
if isComputerUse {
|
||||
providerTools = append(providerTools, chatloop.ProviderTool{
|
||||
Definition: chattool.ComputerUseProviderTool(
|
||||
workspacesdk.DesktopDisplayWidth,
|
||||
workspacesdk.DesktopDisplayHeight),
|
||||
Runner: chattool.NewComputerUseTool(
|
||||
workspacesdk.DesktopDisplayWidth,
|
||||
workspacesdk.DesktopDisplayHeight,
|
||||
getWorkspaceConn, quartz.NewReal(),
|
||||
),
|
||||
})
|
||||
}
|
||||
err = chatloop.Run(ctx, chatloop.RunOptions{
|
||||
Model: model,
|
||||
Messages: prompt,
|
||||
Tools: tools,
|
||||
MaxSteps: maxChatSteps,
|
||||
Tools: tools, MaxSteps: maxChatSteps,
|
||||
|
||||
ModelConfig: callConfig,
|
||||
ProviderOptions: chatprovider.ProviderOptionsFromChatModelConfig(model, callConfig.ProviderOptions),
|
||||
@@ -2668,14 +2702,16 @@ func (p *Server) runChat(
|
||||
// buildProviderTools creates provider-native tool definitions
|
||||
// (like web search) based on the model configuration. These
|
||||
// tools are executed server-side by the LLM provider.
|
||||
func buildProviderTools(_ string, options *codersdk.ChatModelProviderOptions) []fantasy.Tool {
|
||||
var tools []fantasy.Tool
|
||||
func buildProviderTools(_ string, options *codersdk.ChatModelProviderOptions) []chatloop.ProviderTool {
|
||||
var tools []chatloop.ProviderTool
|
||||
|
||||
if options.Anthropic != nil && options.Anthropic.WebSearchEnabled != nil && *options.Anthropic.WebSearchEnabled {
|
||||
tools = append(tools, anthropic.WebSearchTool(&anthropic.WebSearchToolOptions{
|
||||
AllowedDomains: options.Anthropic.AllowedDomains,
|
||||
BlockedDomains: options.Anthropic.BlockedDomains,
|
||||
}))
|
||||
tools = append(tools, chatloop.ProviderTool{
|
||||
Definition: anthropic.WebSearchTool(&anthropic.WebSearchToolOptions{
|
||||
AllowedDomains: options.Anthropic.AllowedDomains,
|
||||
BlockedDomains: options.Anthropic.BlockedDomains,
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
if options.OpenAI != nil && options.OpenAI.WebSearchEnabled != nil && *options.OpenAI.WebSearchEnabled {
|
||||
@@ -2686,17 +2722,21 @@ func buildProviderTools(_ string, options *codersdk.ChatModelProviderOptions) []
|
||||
if len(options.OpenAI.AllowedDomains) > 0 {
|
||||
args["allowed_domains"] = options.OpenAI.AllowedDomains
|
||||
}
|
||||
tools = append(tools, fantasy.ProviderDefinedTool{
|
||||
ID: "web_search",
|
||||
Name: "web_search",
|
||||
Args: args,
|
||||
tools = append(tools, chatloop.ProviderTool{
|
||||
Definition: fantasy.ProviderDefinedTool{
|
||||
ID: "web_search",
|
||||
Name: "web_search",
|
||||
Args: args,
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
if options.Google != nil && options.Google.WebSearchEnabled != nil && *options.Google.WebSearchEnabled {
|
||||
tools = append(tools, fantasy.ProviderDefinedTool{
|
||||
ID: "web_search",
|
||||
Name: "web_search",
|
||||
tools = append(tools, chatloop.ProviderTool{
|
||||
Definition: fantasy.ProviderDefinedTool{
|
||||
ID: "web_search",
|
||||
Name: "web_search",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -6,6 +6,9 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
@@ -14,12 +17,15 @@ import (
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"go.uber.org/mock/gomock"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog/v3/sloggers/slogtest"
|
||||
"github.com/coder/coder/v2/agent/agenttest"
|
||||
"github.com/coder/coder/v2/coderd/chatd"
|
||||
"github.com/coder/coder/v2/coderd/chatd/chatprompt"
|
||||
"github.com/coder/coder/v2/coderd/chatd/chattest"
|
||||
"github.com/coder/coder/v2/coderd/chatd/chattool"
|
||||
"github.com/coder/coder/v2/coderd/coderdtest"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/db2sdk"
|
||||
@@ -28,6 +34,8 @@ import (
|
||||
dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub"
|
||||
"github.com/coder/coder/v2/coderd/util/slice"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock"
|
||||
"github.com/coder/coder/v2/provisioner/echo"
|
||||
proto "github.com/coder/coder/v2/provisionersdk/proto"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
@@ -1767,3 +1775,349 @@ func TestSuccessfulChatSendsWebPushWithSummary(t *testing.T) {
|
||||
require.NotEqual(t, "Agent has finished running.", msg.Body,
|
||||
"push body should not use the default fallback text")
|
||||
}
|
||||
|
||||
func TestComputerUseSubagentToolsAndModel(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
// Track tools and model from the Anthropic LLM calls (the
|
||||
// computer use child chat). We use a raw HTTP handler because
|
||||
// the chattest AnthropicRequest struct does not capture tools.
|
||||
type anthropicCall struct {
|
||||
Model string
|
||||
Tools []string
|
||||
}
|
||||
var anthropicMu sync.Mutex
|
||||
var anthropicCalls []anthropicCall
|
||||
|
||||
anthropicSrv := httptest.NewServer(http.HandlerFunc(
|
||||
func(w http.ResponseWriter, r *http.Request) {
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
var req struct {
|
||||
Model string `json:"model"`
|
||||
Stream bool `json:"stream"`
|
||||
Tools []struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"tools"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
|
||||
names := make([]string, len(req.Tools))
|
||||
for i, tool := range req.Tools {
|
||||
names[i] = tool.Name
|
||||
}
|
||||
anthropicMu.Lock()
|
||||
anthropicCalls = append(anthropicCalls, anthropicCall{
|
||||
Model: req.Model,
|
||||
Tools: names,
|
||||
})
|
||||
anthropicMu.Unlock()
|
||||
|
||||
if !req.Stream {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_ = json.NewEncoder(w).Encode(map[string]any{
|
||||
"id": "msg-test",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": chattool.ComputerUseModelName,
|
||||
"content": []map[string]any{{"type": "text", "text": "Done."}},
|
||||
"stop_reason": "end_turn",
|
||||
"usage": map[string]any{"input_tokens": 10, "output_tokens": 5},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// Stream a minimal Anthropic SSE response.
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
flusher, _ := w.(http.Flusher)
|
||||
|
||||
chunks := []map[string]any{
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": map[string]any{
|
||||
"id": "msg-test",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": chattool.ComputerUseModelName,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": map[string]any{
|
||||
"type": "text",
|
||||
"text": "",
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": 0,
|
||||
"delta": map[string]any{
|
||||
"type": "text_delta",
|
||||
"text": "Done.",
|
||||
},
|
||||
},
|
||||
{"type": "content_block_stop", "index": 0},
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": map[string]any{"stop_reason": "end_turn"},
|
||||
"usage": map[string]any{"output_tokens": 5},
|
||||
},
|
||||
{"type": "message_stop"},
|
||||
}
|
||||
|
||||
for _, chunk := range chunks {
|
||||
chunkBytes, _ := json.Marshal(chunk)
|
||||
eventType, _ := chunk["type"].(string)
|
||||
_, _ = fmt.Fprintf(w, "event: %s\ndata: %s\n\n",
|
||||
eventType, chunkBytes)
|
||||
flusher.Flush()
|
||||
}
|
||||
},
|
||||
))
|
||||
t.Cleanup(anthropicSrv.Close)
|
||||
|
||||
// OpenAI mock for the root chat. The first streaming call
|
||||
// triggers spawn_computer_use_agent; subsequent calls reply
|
||||
// with text.
|
||||
var openAICallCount atomic.Int32
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
if openAICallCount.Add(1) == 1 {
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAIToolCallChunk(
|
||||
"spawn_computer_use_agent",
|
||||
`{"prompt":"do the desktop thing","title":"cu-sub"}`,
|
||||
),
|
||||
)
|
||||
}
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("Done.")...,
|
||||
)
|
||||
})
|
||||
|
||||
// Seed the DB: user, openai-compat provider, model config.
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
_, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{
|
||||
Provider: "openai-compat",
|
||||
DisplayName: "OpenAI Compat",
|
||||
APIKey: "test-key",
|
||||
BaseUrl: openAIURL,
|
||||
CreatedBy: uuid.NullUUID{},
|
||||
Enabled: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
|
||||
Provider: "openai-compat",
|
||||
Model: "gpt-4o-mini",
|
||||
DisplayName: "Test Model",
|
||||
CreatedBy: uuid.NullUUID{},
|
||||
UpdatedBy: uuid.NullUUID{},
|
||||
Enabled: true,
|
||||
IsDefault: true,
|
||||
ContextLimit: 128000,
|
||||
CompressionThreshold: 70,
|
||||
Options: json.RawMessage(`{}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Add an Anthropic provider pointing to our mock server.
|
||||
_, err = db.InsertChatProvider(ctx, database.InsertChatProviderParams{
|
||||
Provider: "anthropic",
|
||||
DisplayName: "Anthropic",
|
||||
APIKey: "test-anthropic-key",
|
||||
BaseUrl: anthropicSrv.URL,
|
||||
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
|
||||
Enabled: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Build workspace + agent records so getWorkspaceConn can
|
||||
// resolve the agent for the computer use child.
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
})
|
||||
tpl := dbgen.Template(t, db, database.Template{
|
||||
CreatedBy: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
ActiveVersionID: tv.ID,
|
||||
})
|
||||
ws := dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
TemplateID: tpl.ID,
|
||||
OwnerID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
|
||||
InitiatorID: user.ID,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
_ = dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
|
||||
TemplateVersionID: tv.ID,
|
||||
WorkspaceID: ws.ID,
|
||||
JobID: pj.ID,
|
||||
})
|
||||
res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{
|
||||
Transition: database.WorkspaceTransitionStart,
|
||||
JobID: pj.ID,
|
||||
})
|
||||
dbAgent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{
|
||||
ResourceID: res.ID,
|
||||
})
|
||||
|
||||
// Mock agent connection that returns valid display dimensions
|
||||
// for the initial screenshot check in the computer use path.
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
mockConn.EXPECT().
|
||||
ExecuteDesktopAction(gomock.Any(), gomock.Any()).
|
||||
Return(workspacesdk.DesktopActionResponse{
|
||||
ScreenshotWidth: 1920,
|
||||
ScreenshotHeight: 1080,
|
||||
ScreenshotData: "iVBOR",
|
||||
}, nil).
|
||||
AnyTimes()
|
||||
mockConn.EXPECT().
|
||||
SetExtraHeaders(gomock.Any()).
|
||||
AnyTimes()
|
||||
mockConn.EXPECT().
|
||||
LS(gomock.Any(), gomock.Any(), gomock.Any()).
|
||||
Return(workspacesdk.LSResponse{}, xerrors.New("not found")).
|
||||
AnyTimes()
|
||||
|
||||
agentConnFn := func(
|
||||
_ context.Context, agentID uuid.UUID,
|
||||
) (workspacesdk.AgentConn, func(), error) {
|
||||
require.Equal(t, dbAgent.ID, agentID)
|
||||
return mockConn, func() {}, nil
|
||||
}
|
||||
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
server := chatd.New(chatd.Config{
|
||||
Logger: logger,
|
||||
Database: db,
|
||||
ReplicaID: uuid.New(),
|
||||
Pubsub: ps,
|
||||
PendingChatAcquireInterval: 10 * time.Millisecond,
|
||||
InFlightChatStaleAfter: testutil.WaitSuperLong,
|
||||
AgentConn: agentConnFn,
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, server.Close())
|
||||
})
|
||||
|
||||
// Create a root chat with a workspace so the child inherits it.
|
||||
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "computer-use-detection",
|
||||
ModelConfigID: model.ID,
|
||||
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText("Use the desktop to check the UI"),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Wait for the root chat AND the computer use child to finish.
|
||||
// The root chat spawns the child, then the chatd server picks
|
||||
// up and runs the child (which hits the Anthropic mock).
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
if got.Status != database.ChatStatusWaiting &&
|
||||
got.Status != database.ChatStatusError {
|
||||
return false
|
||||
}
|
||||
// Ensure the Anthropic mock received at least one call.
|
||||
anthropicMu.Lock()
|
||||
n := len(anthropicCalls)
|
||||
anthropicMu.Unlock()
|
||||
return n >= 1
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
anthropicMu.Lock()
|
||||
calls := append([]anthropicCall(nil), anthropicCalls...)
|
||||
anthropicMu.Unlock()
|
||||
|
||||
require.NotEmpty(t, calls,
|
||||
"expected at least one Anthropic LLM call")
|
||||
|
||||
childModel := calls[0].Model
|
||||
childTools := calls[0].Tools
|
||||
|
||||
// 1. Verify the model is the computer use model.
|
||||
require.Equal(t, chattool.ComputerUseModelName, childModel,
|
||||
"computer use subagent should use %s",
|
||||
chattool.ComputerUseModelName)
|
||||
|
||||
// 2. Verify the computer tool is present.
|
||||
require.Contains(t, childTools, "computer",
|
||||
"computer use subagent should have the computer tool")
|
||||
|
||||
// 3. Verify standard workspace tools are present (the same
|
||||
// set a regular subagent gets).
|
||||
standardTools := []string{
|
||||
"read_file", "write_file", "edit_files", "execute",
|
||||
"process_output", "process_list", "process_signal",
|
||||
}
|
||||
for _, tool := range standardTools {
|
||||
require.Contains(t, childTools, tool,
|
||||
"computer use subagent should have standard tool %q",
|
||||
tool)
|
||||
}
|
||||
|
||||
// 4. Verify workspace provisioning tools are NOT present.
|
||||
workspaceProvisioningTools := []string{
|
||||
"list_templates", "read_template",
|
||||
"create_workspace", "start_workspace",
|
||||
}
|
||||
for _, tool := range workspaceProvisioningTools {
|
||||
require.NotContains(t, childTools, tool,
|
||||
"computer use subagent should NOT have workspace "+
|
||||
"provisioning tool %q", tool)
|
||||
}
|
||||
|
||||
// 5. Verify subagent tools are NOT present.
|
||||
subagentTools := []string{
|
||||
"spawn_agent", "spawn_computer_use_agent",
|
||||
"wait_agent", "message_agent", "close_agent",
|
||||
}
|
||||
for _, tool := range subagentTools {
|
||||
require.NotContains(t, childTools, tool,
|
||||
"computer use subagent should NOT have subagent "+
|
||||
"tool %q", tool)
|
||||
}
|
||||
|
||||
// 6. Verify the child chat has Mode = computer_use in
|
||||
// the DB.
|
||||
allChats, err := db.GetChatsByOwnerID(ctx, database.GetChatsByOwnerIDParams{
|
||||
OwnerID: user.ID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
var children []database.Chat
|
||||
for _, c := range allChats {
|
||||
if c.ParentChatID.Valid && c.ParentChatID.UUID == chat.ID {
|
||||
children = append(children, c)
|
||||
}
|
||||
}
|
||||
require.Len(t, children, 1)
|
||||
require.True(t, children[0].Mode.Valid)
|
||||
require.Equal(t, database.ChatModeComputerUse,
|
||||
children[0].Mode.ChatMode)
|
||||
}
|
||||
|
||||
@@ -63,11 +63,12 @@ type RunOptions struct {
|
||||
// of the provider, which lives in chatd, not chatloop.
|
||||
ProviderOptions fantasy.ProviderOptions
|
||||
|
||||
// ProviderTools are provider-native tools (like web search)
|
||||
// that are passed directly to the provider API alongside
|
||||
// function tool definitions. These are not necessarily
|
||||
// executed server-side; handling is provider-specific.
|
||||
ProviderTools []fantasy.Tool
|
||||
// ProviderTools are provider-native tools (like web search
|
||||
// and computer use) whose definitions are passed directly
|
||||
// to the provider API. When a ProviderTool has a non-nil
|
||||
// Runner, tool calls are executed locally; otherwise the
|
||||
// provider handles execution (e.g. web search).
|
||||
ProviderTools []ProviderTool
|
||||
|
||||
PersistStep func(context.Context, PersistedStep) error
|
||||
PublishMessagePart func(
|
||||
@@ -88,6 +89,16 @@ type RunOptions struct {
|
||||
OnInterruptedPersistError func(error)
|
||||
}
|
||||
|
||||
// ProviderTool pairs a provider-native tool definition with an
|
||||
// optional local executor. When Runner is nil the tool is fully
|
||||
// provider-executed (e.g. web search). When Runner is non-nil
|
||||
// the definition is sent to the API but execution is handled
|
||||
// locally (e.g. computer use).
|
||||
type ProviderTool struct {
|
||||
Definition fantasy.Tool
|
||||
Runner fantasy.AgentTool
|
||||
}
|
||||
|
||||
// stepResult holds the accumulated output of a single streaming
|
||||
// step. Since we own the stream consumer, all content is tracked
|
||||
// directly here — no shadow draft state needed.
|
||||
@@ -315,7 +326,7 @@ func Run(ctx context.Context, opts RunOptions) error {
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
toolResults = executeTools(ctx, opts.Tools, result.toolCalls, func(tr fantasy.ToolResultContent) {
|
||||
toolResults = executeTools(ctx, opts.Tools, opts.ProviderTools, result.toolCalls, func(tr fantasy.ToolResultContent) {
|
||||
publishMessagePart(
|
||||
codersdk.ChatMessageRoleTool,
|
||||
chatprompt.PartFromContent(tr),
|
||||
@@ -639,6 +650,7 @@ func processStepStream(
|
||||
func executeTools(
|
||||
ctx context.Context,
|
||||
allTools []fantasy.AgentTool,
|
||||
providerTools []ProviderTool,
|
||||
toolCalls []fantasy.ToolCallContent,
|
||||
onResult func(fantasy.ToolResultContent),
|
||||
) []fantasy.ToolResultContent {
|
||||
@@ -664,6 +676,13 @@ func executeTools(
|
||||
for _, t := range allTools {
|
||||
toolMap[t.Info().Name] = t
|
||||
}
|
||||
// Include runners from provider tools so locally-executed
|
||||
// provider tools (e.g. computer use) can be dispatched.
|
||||
for _, pt := range providerTools {
|
||||
if pt.Runner != nil {
|
||||
toolMap[pt.Runner.Info().Name] = pt.Runner
|
||||
}
|
||||
}
|
||||
|
||||
results := make([]fantasy.ToolResultContent, len(localToolCalls))
|
||||
var wg sync.WaitGroup
|
||||
@@ -863,15 +882,16 @@ func persistInterruptedStep(
|
||||
// buildToolDefinitions converts AgentTool definitions into the
|
||||
// fantasy.Tool slice expected by fantasy.Call. When activeTools
|
||||
// is non-empty, only function tools whose name appears in the
|
||||
// list are included. Provider tools bypass this filter and are
|
||||
// always appended unconditionally.
|
||||
func buildToolDefinitions(tools []fantasy.AgentTool, activeTools []string, providerTools []fantasy.Tool) []fantasy.Tool {
|
||||
prepared := make([]fantasy.Tool, 0, len(tools))
|
||||
// list are included. Provider tool definitions are always
|
||||
// appended unconditionally.
|
||||
func buildToolDefinitions(tools []fantasy.AgentTool, activeTools []string, providerTools []ProviderTool) []fantasy.Tool {
|
||||
prepared := make([]fantasy.Tool, 0, len(tools)+len(providerTools))
|
||||
for _, tool := range tools {
|
||||
info := tool.Info()
|
||||
if len(activeTools) > 0 && !slices.Contains(activeTools, info.Name) {
|
||||
continue
|
||||
}
|
||||
|
||||
inputSchema := map[string]any{
|
||||
"type": "object",
|
||||
"properties": info.Parameters,
|
||||
@@ -885,7 +905,9 @@ func buildToolDefinitions(tools []fantasy.AgentTool, activeTools []string, provi
|
||||
ProviderOptions: tool.ProviderOptions(),
|
||||
})
|
||||
}
|
||||
prepared = append(prepared, providerTools...)
|
||||
for _, pt := range providerTools {
|
||||
prepared = append(prepared, pt.Definition)
|
||||
}
|
||||
return prepared
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
package chattool
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"time"
|
||||
|
||||
"charm.land/fantasy"
|
||||
fantasyanthropic "charm.land/fantasy/providers/anthropic"
|
||||
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
const (
|
||||
// ComputerUseModelProvider is the provider for the computer
|
||||
// use model.
|
||||
ComputerUseModelProvider = "anthropic"
|
||||
// ComputerUseModelName is the model used for computer use
|
||||
// subagents.
|
||||
ComputerUseModelName = "claude-opus-4-6"
|
||||
)
|
||||
|
||||
// computerUseTool implements fantasy.AgentTool and
|
||||
// chatloop.ToolDefiner for Anthropic computer use.
|
||||
type computerUseTool struct {
|
||||
displayWidth int
|
||||
displayHeight int
|
||||
getWorkspaceConn func(ctx context.Context) (workspacesdk.AgentConn, error)
|
||||
providerOptions fantasy.ProviderOptions
|
||||
clock quartz.Clock
|
||||
}
|
||||
|
||||
// NewComputerUseTool creates a computer use AgentTool that
|
||||
// delegates to the agent's desktop endpoints.
|
||||
func NewComputerUseTool(
|
||||
displayWidth, displayHeight int,
|
||||
getWorkspaceConn func(ctx context.Context) (workspacesdk.AgentConn, error),
|
||||
clock quartz.Clock,
|
||||
) fantasy.AgentTool {
|
||||
return &computerUseTool{
|
||||
displayWidth: displayWidth,
|
||||
displayHeight: displayHeight,
|
||||
getWorkspaceConn: getWorkspaceConn,
|
||||
clock: clock,
|
||||
}
|
||||
}
|
||||
|
||||
func (*computerUseTool) Info() fantasy.ToolInfo {
|
||||
return fantasy.ToolInfo{
|
||||
Name: "computer",
|
||||
Description: "Control the desktop: take screenshots, move the mouse, click, type, and scroll.",
|
||||
Parameters: map[string]any{},
|
||||
Required: []string{},
|
||||
}
|
||||
}
|
||||
|
||||
// ComputerUseProviderTool creates the provider-defined tool
|
||||
// definition for Anthropic computer use. This is passed via
|
||||
// ProviderTools so the API receives the correct wire format.
|
||||
func ComputerUseProviderTool(displayWidth, displayHeight int) fantasy.Tool {
|
||||
return fantasyanthropic.NewComputerUseTool(
|
||||
fantasyanthropic.ComputerUseToolOptions{
|
||||
DisplayWidthPx: int64(displayWidth),
|
||||
DisplayHeightPx: int64(displayHeight),
|
||||
ToolVersion: fantasyanthropic.ComputerUse20251124,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func (t *computerUseTool) ProviderOptions() fantasy.ProviderOptions {
|
||||
return t.providerOptions
|
||||
}
|
||||
|
||||
func (t *computerUseTool) SetProviderOptions(opts fantasy.ProviderOptions) {
|
||||
t.providerOptions = opts
|
||||
}
|
||||
|
||||
func (t *computerUseTool) Run(ctx context.Context, call fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
input, err := fantasyanthropic.ParseComputerUseInput(call.Input)
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
fmt.Sprintf("invalid computer use input: %v", err),
|
||||
), nil
|
||||
}
|
||||
|
||||
conn, err := t.getWorkspaceConn(ctx)
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
fmt.Sprintf("failed to connect to workspace: %v", err),
|
||||
), nil
|
||||
}
|
||||
|
||||
// Compute scaled screenshot size for Anthropic constraints.
|
||||
scaledW, scaledH := computeScaledScreenshotSize(
|
||||
t.displayWidth, t.displayHeight,
|
||||
)
|
||||
|
||||
// For wait actions, sleep then return a screenshot.
|
||||
if input.Action == fantasyanthropic.ActionWait {
|
||||
d := input.Duration
|
||||
if d <= 0 {
|
||||
d = 1000
|
||||
}
|
||||
timer := t.clock.NewTimer(time.Duration(d)*time.Millisecond, "computeruse", "wait")
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
case <-timer.C:
|
||||
}
|
||||
screenshotAction := workspacesdk.DesktopAction{
|
||||
Action: "screenshot",
|
||||
ScaledWidth: &scaledW,
|
||||
ScaledHeight: &scaledH,
|
||||
}
|
||||
screenResp, sErr := conn.ExecuteDesktopAction(ctx, screenshotAction)
|
||||
if sErr != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
fmt.Sprintf("screenshot failed: %v", sErr),
|
||||
), nil
|
||||
}
|
||||
return fantasy.NewImageResponse(
|
||||
[]byte(screenResp.ScreenshotData), "image/png",
|
||||
), nil
|
||||
}
|
||||
|
||||
// For screenshot action, use ExecuteDesktopAction.
|
||||
if input.Action == fantasyanthropic.ActionScreenshot {
|
||||
screenshotAction := workspacesdk.DesktopAction{
|
||||
Action: "screenshot",
|
||||
ScaledWidth: &scaledW,
|
||||
ScaledHeight: &scaledH,
|
||||
}
|
||||
screenResp, sErr := conn.ExecuteDesktopAction(ctx, screenshotAction)
|
||||
if sErr != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
fmt.Sprintf("screenshot failed: %v", sErr),
|
||||
), nil
|
||||
}
|
||||
return fantasy.NewImageResponse(
|
||||
[]byte(screenResp.ScreenshotData), "image/png",
|
||||
), nil
|
||||
}
|
||||
|
||||
// Build the action request.
|
||||
action := workspacesdk.DesktopAction{
|
||||
Action: string(input.Action),
|
||||
ScaledWidth: &scaledW,
|
||||
ScaledHeight: &scaledH,
|
||||
}
|
||||
if input.Coordinate != ([2]int64{}) {
|
||||
coord := [2]int{int(input.Coordinate[0]), int(input.Coordinate[1])}
|
||||
action.Coordinate = &coord
|
||||
}
|
||||
if input.StartCoordinate != ([2]int64{}) {
|
||||
coord := [2]int{int(input.StartCoordinate[0]), int(input.StartCoordinate[1])}
|
||||
action.StartCoordinate = &coord
|
||||
}
|
||||
if input.Text != "" {
|
||||
action.Text = &input.Text
|
||||
}
|
||||
if input.Duration > 0 {
|
||||
d := int(input.Duration)
|
||||
action.Duration = &d
|
||||
}
|
||||
if input.ScrollAmount > 0 {
|
||||
s := int(input.ScrollAmount)
|
||||
action.ScrollAmount = &s
|
||||
}
|
||||
if input.ScrollDirection != "" {
|
||||
action.ScrollDirection = &input.ScrollDirection
|
||||
}
|
||||
|
||||
// Execute the action.
|
||||
_, err = conn.ExecuteDesktopAction(ctx, action)
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
fmt.Sprintf("action %q failed: %v", input.Action, err),
|
||||
), nil
|
||||
}
|
||||
|
||||
// Take a screenshot after every action (Anthropic pattern).
|
||||
screenshotAction := workspacesdk.DesktopAction{
|
||||
Action: "screenshot",
|
||||
ScaledWidth: &scaledW,
|
||||
ScaledHeight: &scaledH,
|
||||
}
|
||||
screenResp, sErr := conn.ExecuteDesktopAction(ctx, screenshotAction)
|
||||
if sErr != nil {
|
||||
return fantasy.NewTextErrorResponse(
|
||||
fmt.Sprintf("screenshot failed: %v", sErr),
|
||||
), nil
|
||||
}
|
||||
|
||||
return fantasy.NewImageResponse(
|
||||
[]byte(screenResp.ScreenshotData), "image/png",
|
||||
), nil
|
||||
}
|
||||
|
||||
// computeScaledScreenshotSize computes the target screenshot
|
||||
// dimensions to fit within Anthropic's constraints.
|
||||
func computeScaledScreenshotSize(width, height int) (scaledWidth int, scaledHeight int) {
|
||||
const maxLongEdge = 1568
|
||||
const maxTotalPixels = 1_150_000
|
||||
|
||||
longEdge := max(width, height)
|
||||
totalPixels := width * height
|
||||
longEdgeScale := float64(maxLongEdge) / float64(longEdge)
|
||||
totalPixelsScale := math.Sqrt(
|
||||
float64(maxTotalPixels) / float64(totalPixels),
|
||||
)
|
||||
scale := min(1.0, longEdgeScale, totalPixelsScale)
|
||||
|
||||
if scale >= 1.0 {
|
||||
return width, height
|
||||
}
|
||||
return max(1, int(float64(width)*scale)),
|
||||
max(1, int(float64(height)*scale))
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
package chattool
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestComputeScaledScreenshotSize(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
width, height int
|
||||
wantW, wantH int
|
||||
}{
|
||||
{
|
||||
name: "1920x1080_scales_down",
|
||||
width: 1920,
|
||||
height: 1080,
|
||||
wantW: 1429,
|
||||
wantH: 804,
|
||||
},
|
||||
{
|
||||
name: "1280x800_no_scaling",
|
||||
width: 1280,
|
||||
height: 800,
|
||||
wantW: 1280,
|
||||
wantH: 800,
|
||||
},
|
||||
{
|
||||
name: "3840x2160_large_display",
|
||||
width: 3840,
|
||||
height: 2160,
|
||||
wantW: 1429,
|
||||
wantH: 804,
|
||||
},
|
||||
{
|
||||
name: "1568x1000_pixel_cap_applies",
|
||||
width: 1568,
|
||||
height: 1000,
|
||||
wantW: 1342,
|
||||
wantH: 856,
|
||||
},
|
||||
{
|
||||
name: "100x100_small_display",
|
||||
width: 100,
|
||||
height: 100,
|
||||
wantW: 100,
|
||||
wantH: 100,
|
||||
},
|
||||
{
|
||||
name: "4000x3000_stays_within_limits",
|
||||
width: 4000,
|
||||
// Both constraints apply. The function should keep
|
||||
// the result within maxLongEdge=1568 and
|
||||
// totalPixels<=1,150,000.
|
||||
height: 3000,
|
||||
wantW: 1238,
|
||||
wantH: 928,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
gotW, gotH := computeScaledScreenshotSize(tt.width, tt.height)
|
||||
assert.Equal(t, tt.wantW, gotW)
|
||||
assert.Equal(t, tt.wantH, gotH)
|
||||
|
||||
// Invariant: results must respect Anthropic constraints.
|
||||
const maxLongEdge = 1568
|
||||
const maxTotalPixels = 1_150_000
|
||||
longEdge := max(gotW, gotH)
|
||||
assert.LessOrEqual(t, longEdge, maxLongEdge,
|
||||
"long edge %d exceeds max %d", longEdge, maxLongEdge)
|
||||
assert.LessOrEqual(t, gotW*gotH, maxTotalPixels,
|
||||
"total pixels %d exceeds max %d", gotW*gotH, maxTotalPixels)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
package chattool_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"charm.land/fantasy"
|
||||
"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/chatd/chattool"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk"
|
||||
"github.com/coder/coder/v2/codersdk/workspacesdk/agentconnmock"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
func TestComputerUseTool_Info(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := chattool.NewComputerUseTool(workspacesdk.DesktopDisplayWidth, workspacesdk.DesktopDisplayHeight, nil, quartz.NewReal())
|
||||
info := tool.Info()
|
||||
assert.Equal(t, "computer", info.Name)
|
||||
assert.NotEmpty(t, info.Description)
|
||||
}
|
||||
|
||||
func TestComputerUseProviderTool(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
def := chattool.ComputerUseProviderTool(workspacesdk.DesktopDisplayWidth, workspacesdk.DesktopDisplayHeight)
|
||||
pdt, ok := def.(fantasy.ProviderDefinedTool)
|
||||
require.True(t, ok, "ComputerUseProviderTool should return a ProviderDefinedTool")
|
||||
assert.Contains(t, pdt.ID, "computer")
|
||||
assert.Equal(t, "computer", pdt.Name)
|
||||
// Verify display dimensions are passed through.
|
||||
assert.Equal(t, int64(workspacesdk.DesktopDisplayWidth), pdt.Args["display_width_px"])
|
||||
assert.Equal(t, int64(workspacesdk.DesktopDisplayHeight), pdt.Args["display_height_px"])
|
||||
}
|
||||
|
||||
func TestComputerUseTool_Run_Screenshot(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
|
||||
mockConn.EXPECT().ExecuteDesktopAction(
|
||||
gomock.Any(),
|
||||
gomock.Any(),
|
||||
).Return(workspacesdk.DesktopActionResponse{
|
||||
Output: "screenshot",
|
||||
ScreenshotData: "base64png",
|
||||
ScreenshotWidth: 1024,
|
||||
ScreenshotHeight: 768,
|
||||
}, nil)
|
||||
|
||||
tool := chattool.NewComputerUseTool(workspacesdk.DesktopDisplayWidth, workspacesdk.DesktopDisplayHeight, func(_ context.Context) (workspacesdk.AgentConn, error) {
|
||||
return mockConn, nil
|
||||
}, quartz.NewReal())
|
||||
|
||||
call := fantasy.ToolCall{
|
||||
ID: "test-1",
|
||||
Name: "computer",
|
||||
Input: `{"action":"screenshot"}`,
|
||||
}
|
||||
|
||||
resp, err := tool.Run(context.Background(), call)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "image", resp.Type)
|
||||
assert.Equal(t, "image/png", resp.MediaType)
|
||||
assert.Equal(t, []byte("base64png"), resp.Data)
|
||||
assert.False(t, resp.IsError)
|
||||
}
|
||||
|
||||
func TestComputerUseTool_Run_LeftClick(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
|
||||
// Expect the action call first.
|
||||
mockConn.EXPECT().ExecuteDesktopAction(
|
||||
gomock.Any(),
|
||||
gomock.Any(),
|
||||
).Return(workspacesdk.DesktopActionResponse{
|
||||
Output: "left_click performed",
|
||||
}, nil)
|
||||
|
||||
// Then expect a screenshot (auto-screenshot after action).
|
||||
mockConn.EXPECT().ExecuteDesktopAction(
|
||||
gomock.Any(),
|
||||
gomock.Any(),
|
||||
).Return(workspacesdk.DesktopActionResponse{
|
||||
Output: "screenshot",
|
||||
ScreenshotData: "after-click",
|
||||
ScreenshotWidth: 1024,
|
||||
ScreenshotHeight: 768,
|
||||
}, nil)
|
||||
|
||||
tool := chattool.NewComputerUseTool(workspacesdk.DesktopDisplayWidth, workspacesdk.DesktopDisplayHeight, func(_ context.Context) (workspacesdk.AgentConn, error) {
|
||||
return mockConn, nil
|
||||
}, quartz.NewReal())
|
||||
|
||||
call := fantasy.ToolCall{
|
||||
ID: "test-2",
|
||||
Name: "computer",
|
||||
Input: `{"action":"left_click","coordinate":[100,200]}`,
|
||||
}
|
||||
|
||||
resp, err := tool.Run(context.Background(), call)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "image", resp.Type)
|
||||
assert.Equal(t, []byte("after-click"), resp.Data)
|
||||
}
|
||||
|
||||
func TestComputerUseTool_Run_Wait(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctrl := gomock.NewController(t)
|
||||
mockConn := agentconnmock.NewMockAgentConn(ctrl)
|
||||
// Expect a screenshot after the wait completes.
|
||||
mockConn.EXPECT().ExecuteDesktopAction(
|
||||
gomock.Any(),
|
||||
gomock.Any(),
|
||||
).Return(workspacesdk.DesktopActionResponse{
|
||||
Output: "screenshot",
|
||||
ScreenshotData: "after-wait",
|
||||
ScreenshotWidth: 1024,
|
||||
ScreenshotHeight: 768,
|
||||
}, nil)
|
||||
|
||||
tool := chattool.NewComputerUseTool(workspacesdk.DesktopDisplayWidth, workspacesdk.DesktopDisplayHeight, func(_ context.Context) (workspacesdk.AgentConn, error) {
|
||||
return mockConn, nil
|
||||
}, quartz.NewReal())
|
||||
|
||||
call := fantasy.ToolCall{
|
||||
ID: "test-3",
|
||||
Name: "computer",
|
||||
Input: `{"action":"wait","duration":10}`,
|
||||
}
|
||||
|
||||
resp, err := tool.Run(context.Background(), call)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "image", resp.Type)
|
||||
assert.Equal(t, "image/png", resp.MediaType)
|
||||
assert.Equal(t, []byte("after-wait"), resp.Data)
|
||||
assert.False(t, resp.IsError)
|
||||
}
|
||||
|
||||
func TestComputerUseTool_Run_ConnError(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := chattool.NewComputerUseTool(workspacesdk.DesktopDisplayWidth, workspacesdk.DesktopDisplayHeight, func(_ context.Context) (workspacesdk.AgentConn, error) {
|
||||
return nil, xerrors.New("workspace not available")
|
||||
}, quartz.NewReal())
|
||||
|
||||
call := fantasy.ToolCall{
|
||||
ID: "test-4",
|
||||
Name: "computer",
|
||||
Input: `{"action":"screenshot"}`,
|
||||
}
|
||||
|
||||
resp, err := tool.Run(context.Background(), call)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, resp.IsError)
|
||||
assert.Contains(t, resp.Content, "workspace not available")
|
||||
}
|
||||
|
||||
func TestComputerUseTool_Run_InvalidInput(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tool := chattool.NewComputerUseTool(workspacesdk.DesktopDisplayWidth, workspacesdk.DesktopDisplayHeight, func(_ context.Context) (workspacesdk.AgentConn, error) {
|
||||
return nil, xerrors.New("should not be called")
|
||||
}, quartz.NewReal())
|
||||
|
||||
call := fantasy.ToolCall{
|
||||
ID: "test-5",
|
||||
Name: "computer",
|
||||
Input: `{invalid json`,
|
||||
}
|
||||
|
||||
resp, err := tool.Run(context.Background(), call)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, resp.IsError)
|
||||
assert.Contains(t, resp.Content, "invalid computer use input")
|
||||
}
|
||||
+123
-2
@@ -13,6 +13,7 @@ import (
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/chatd/chatprompt"
|
||||
"github.com/coder/coder/v2/coderd/chatd/chatprovider"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
coderdpubsub "github.com/coder/coder/v2/coderd/pubsub"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
@@ -26,11 +27,30 @@ const (
|
||||
defaultSubagentWaitTimeout = 5 * time.Minute
|
||||
)
|
||||
|
||||
// computerUseSubagentSystemPrompt is the system prompt prepended to
|
||||
// every computer use subagent chat. It instructs the model on how to
|
||||
// interact with the desktop environment via the computer tool.
|
||||
const computerUseSubagentSystemPrompt = `You are a computer use agent with access to a desktop environment. You can see the screen, move the mouse, click, type, scroll, and drag.
|
||||
|
||||
Your primary tool is the "computer" tool which lets you interact with the desktop. After every action you take, you will receive a screenshot showing the current state of the screen. Use these screenshots to verify your actions and plan next steps.
|
||||
|
||||
Guidelines:
|
||||
- Always start by taking a screenshot to see the current state of the desktop.
|
||||
- Be precise with coordinates when clicking or typing.
|
||||
- Wait for UI elements to load before interacting with them.
|
||||
- If an action doesn't produce the expected result, try alternative approaches.
|
||||
- Report what you accomplished when done.`
|
||||
|
||||
type spawnAgentArgs struct {
|
||||
Prompt string `json:"prompt"`
|
||||
Title string `json:"title,omitempty"`
|
||||
}
|
||||
|
||||
type spawnComputerUseAgentArgs struct {
|
||||
Prompt string `json:"prompt"`
|
||||
Title string `json:"title,omitempty"`
|
||||
}
|
||||
|
||||
type waitAgentArgs struct {
|
||||
ChatID string `json:"chat_id"`
|
||||
TimeoutSeconds *int `json:"timeout_seconds,omitempty"`
|
||||
@@ -46,8 +66,26 @@ type closeAgentArgs struct {
|
||||
ChatID string `json:"chat_id"`
|
||||
}
|
||||
|
||||
func (p *Server) subagentTools(currentChat func() database.Chat) []fantasy.AgentTool {
|
||||
return []fantasy.AgentTool{
|
||||
// isAnthropicConfigured reports whether an Anthropic API key is
|
||||
// available, either from static provider keys or from the database.
|
||||
func (p *Server) isAnthropicConfigured(ctx context.Context) bool {
|
||||
if p.providerAPIKeys.APIKey("anthropic") != "" {
|
||||
return true
|
||||
}
|
||||
dbProviders, err := p.db.GetEnabledChatProviders(ctx)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
for _, prov := range dbProviders {
|
||||
if chatprovider.NormalizeProvider(prov.Provider) == "anthropic" && strings.TrimSpace(prov.APIKey) != "" {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (p *Server) subagentTools(ctx context.Context, currentChat func() database.Chat) []fantasy.AgentTool {
|
||||
tools := []fantasy.AgentTool{
|
||||
fantasy.NewAgentTool(
|
||||
"spawn_agent",
|
||||
"Spawn a delegated child agent to work on a clearly scoped, "+
|
||||
@@ -213,6 +251,89 @@ func (p *Server) subagentTools(currentChat func() database.Chat) []fantasy.Agent
|
||||
},
|
||||
),
|
||||
}
|
||||
|
||||
// Only include the computer use tool when an Anthropic
|
||||
// provider is configured, since it requires an Anthropic
|
||||
// model.
|
||||
if p.isAnthropicConfigured(ctx) {
|
||||
tools = append(tools, fantasy.NewAgentTool(
|
||||
"spawn_computer_use_agent",
|
||||
"Spawn a dedicated computer use agent that can see the desktop "+
|
||||
"(take screenshots) and interact with it (mouse, keyboard, "+
|
||||
"scroll). The agent runs on a model optimized for computer "+
|
||||
"use and has the same workspace tools as a standard subagent "+
|
||||
"plus the native Anthropic computer tool. Use this for tasks "+
|
||||
"that require visual interaction with a desktop GUI (e.g. "+
|
||||
"browser automation, GUI testing, visual inspection). After "+
|
||||
"spawning, use wait_agent to collect the result.",
|
||||
func(ctx context.Context, args spawnComputerUseAgentArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
if currentChat == nil {
|
||||
return fantasy.NewTextErrorResponse("subagent callbacks are not configured"), nil
|
||||
}
|
||||
|
||||
parent := currentChat()
|
||||
if parent.ParentChatID.Valid {
|
||||
return fantasy.NewTextErrorResponse("delegated chats cannot create child subagents"), nil
|
||||
}
|
||||
|
||||
parent, err := p.db.GetChatByID(ctx, parent.ID)
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(err.Error()), nil
|
||||
}
|
||||
|
||||
prompt := strings.TrimSpace(args.Prompt)
|
||||
if prompt == "" {
|
||||
return fantasy.NewTextErrorResponse("prompt is required"), nil
|
||||
}
|
||||
|
||||
title := strings.TrimSpace(args.Title)
|
||||
if title == "" {
|
||||
title = subagentFallbackChatTitle(prompt)
|
||||
}
|
||||
|
||||
rootChatID := parent.ID
|
||||
if parent.RootChatID.Valid {
|
||||
rootChatID = parent.RootChatID.UUID
|
||||
}
|
||||
if parent.LastModelConfigID == uuid.Nil {
|
||||
return fantasy.NewTextErrorResponse("parent chat model config id is required"), nil
|
||||
}
|
||||
|
||||
// Create the child chat with Mode set to
|
||||
// computer_use. This signals runChat to use the
|
||||
// predefined computer use model and include the
|
||||
// computer tool.
|
||||
childChat, err := p.CreateChat(ctx, CreateOptions{
|
||||
OwnerID: parent.OwnerID,
|
||||
WorkspaceID: parent.WorkspaceID,
|
||||
ParentChatID: uuid.NullUUID{
|
||||
UUID: parent.ID,
|
||||
Valid: true,
|
||||
},
|
||||
RootChatID: uuid.NullUUID{
|
||||
UUID: rootChatID,
|
||||
Valid: true,
|
||||
},
|
||||
ModelConfigID: parent.LastModelConfigID,
|
||||
Title: title,
|
||||
ChatMode: database.NullChatMode{ChatMode: database.ChatModeComputerUse, Valid: true},
|
||||
SystemPrompt: computerUseSubagentSystemPrompt + "\n\n" + prompt,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText(prompt)},
|
||||
})
|
||||
if err != nil {
|
||||
return fantasy.NewTextErrorResponse(err.Error()), nil
|
||||
}
|
||||
|
||||
return toolJSONResponse(map[string]any{
|
||||
"chat_id": childChat.ID.String(),
|
||||
"title": childChat.Title,
|
||||
"status": string(childChat.Status),
|
||||
}), nil
|
||||
},
|
||||
))
|
||||
}
|
||||
|
||||
return tools
|
||||
}
|
||||
|
||||
func parseSubagentToolChatID(raw string) (uuid.UUID, error) {
|
||||
|
||||
@@ -0,0 +1,300 @@
|
||||
package chatd
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"charm.land/fantasy"
|
||||
"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/coderd/chatd/chatprovider"
|
||||
"github.com/coder/coder/v2/coderd/chatd/chattool"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbgen"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/coderd/database/pubsub"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
func TestComputerUseSubagentSystemPrompt(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Verify the system prompt constant is non-empty and contains
|
||||
// key instructions for the computer use agent.
|
||||
assert.NotEmpty(t, computerUseSubagentSystemPrompt)
|
||||
assert.Contains(t, computerUseSubagentSystemPrompt, "computer")
|
||||
assert.Contains(t, computerUseSubagentSystemPrompt, "screenshot")
|
||||
}
|
||||
|
||||
func TestSubagentFallbackChatTitle(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
input string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "EmptyPrompt",
|
||||
input: "",
|
||||
want: "New Chat",
|
||||
},
|
||||
{
|
||||
name: "ShortPrompt",
|
||||
input: "Open Firefox",
|
||||
want: "Open Firefox",
|
||||
},
|
||||
{
|
||||
name: "LongPrompt",
|
||||
input: "Please open the Firefox browser and navigate to the settings page",
|
||||
want: "Please open the Firefox browser and...",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
got := subagentFallbackChatTitle(tt.input)
|
||||
assert.Equal(t, tt.want, got)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// newInternalTestServer creates a Server for internal tests with
|
||||
// custom provider API keys. The server is automatically closed
|
||||
// when the test finishes.
|
||||
func newInternalTestServer(
|
||||
t *testing.T,
|
||||
db database.Store,
|
||||
ps pubsub.Pubsub,
|
||||
keys chatprovider.ProviderAPIKeys,
|
||||
) *Server {
|
||||
t.Helper()
|
||||
|
||||
logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})
|
||||
server := New(Config{
|
||||
Logger: logger,
|
||||
Database: db,
|
||||
ReplicaID: uuid.New(),
|
||||
Pubsub: ps,
|
||||
// Use a very long interval so the background loop
|
||||
// does not interfere with test assertions.
|
||||
PendingChatAcquireInterval: testutil.WaitLong,
|
||||
ProviderAPIKeys: keys,
|
||||
})
|
||||
t.Cleanup(func() {
|
||||
require.NoError(t, server.Close())
|
||||
})
|
||||
return server
|
||||
}
|
||||
|
||||
// seedInternalChatDeps inserts an OpenAI provider and model config
|
||||
// into the database and returns the created user and model. This
|
||||
// deliberately does NOT create an Anthropic provider.
|
||||
func seedInternalChatDeps(
|
||||
ctx context.Context,
|
||||
t *testing.T,
|
||||
db database.Store,
|
||||
) (database.User, database.ChatModelConfig) {
|
||||
t.Helper()
|
||||
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
_, err := db.InsertChatProvider(ctx, database.InsertChatProviderParams{
|
||||
Provider: "openai",
|
||||
DisplayName: "OpenAI",
|
||||
APIKey: "test-key",
|
||||
BaseUrl: "",
|
||||
ApiKeyKeyID: sql.NullString{},
|
||||
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
|
||||
Enabled: true,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
model, err := db.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{
|
||||
Provider: "openai",
|
||||
Model: "gpt-4o-mini",
|
||||
DisplayName: "Test Model",
|
||||
CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
|
||||
UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true},
|
||||
Enabled: true,
|
||||
IsDefault: true,
|
||||
ContextLimit: 128000,
|
||||
CompressionThreshold: 70,
|
||||
Options: json.RawMessage(`{}`),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
return user, model
|
||||
}
|
||||
|
||||
// findToolByName returns the tool with the given name from the
|
||||
// slice, or nil if no match is found.
|
||||
func findToolByName(tools []fantasy.AgentTool, name string) fantasy.AgentTool {
|
||||
for _, tool := range tools {
|
||||
if tool.Info().Name == name {
|
||||
return tool
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestSpawnComputerUseAgent_NoAnthropicProvider(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
// No Anthropic key in ProviderAPIKeys.
|
||||
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{})
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
user, model := seedInternalChatDeps(ctx, t, db)
|
||||
|
||||
// Create a root parent chat.
|
||||
parent, err := server.CreateChat(ctx, CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "parent-no-anthropic",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Re-fetch so LastModelConfigID is populated from the DB.
|
||||
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
tools := server.subagentTools(ctx, func() database.Chat { return parentChat })
|
||||
tool := findToolByName(tools, "spawn_computer_use_agent")
|
||||
assert.Nil(t, tool, "spawn_computer_use_agent tool must be omitted when Anthropic is not configured")
|
||||
}
|
||||
|
||||
func TestSpawnComputerUseAgent_NotAvailableForChildChats(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
// Provide an Anthropic key so the provider check passes.
|
||||
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{
|
||||
Anthropic: "test-anthropic-key",
|
||||
})
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
user, model := seedInternalChatDeps(ctx, t, db)
|
||||
|
||||
// Create a root parent chat.
|
||||
parent, err := server.CreateChat(ctx, CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "root-parent",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Create a child chat under the parent.
|
||||
child, err := server.CreateChat(ctx, CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
ParentChatID: uuid.NullUUID{
|
||||
UUID: parent.ID,
|
||||
Valid: true,
|
||||
},
|
||||
RootChatID: uuid.NullUUID{
|
||||
UUID: parent.ID,
|
||||
Valid: true,
|
||||
},
|
||||
Title: "child-subagent",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("do something")},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Re-fetch the child so ParentChatID is populated.
|
||||
childChat, err := db.GetChatByID(ctx, child.ID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, childChat.ParentChatID.Valid,
|
||||
"child chat must have a parent")
|
||||
|
||||
// Get tools as if the child chat is the current chat.
|
||||
tools := server.subagentTools(ctx, func() database.Chat { return childChat })
|
||||
tool := findToolByName(tools, "spawn_computer_use_agent")
|
||||
require.NotNil(t, tool, "spawn_computer_use_agent tool must be present")
|
||||
|
||||
resp, err := tool.Run(ctx, fantasy.ToolCall{
|
||||
ID: "call-2",
|
||||
Name: "spawn_computer_use_agent",
|
||||
Input: `{"prompt":"open browser"}`,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.True(t, resp.IsError, "expected an error response")
|
||||
assert.Contains(t, resp.Content, "delegated chats cannot create child subagents")
|
||||
}
|
||||
|
||||
func TestSpawnComputerUseAgent_UsesComputerUseModelNotParent(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
// Provide an Anthropic key so the tool can proceed.
|
||||
server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{
|
||||
Anthropic: "test-anthropic-key",
|
||||
})
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
user, model := seedInternalChatDeps(ctx, t, db)
|
||||
|
||||
// The parent uses an OpenAI model.
|
||||
require.Equal(t, "openai", model.Provider,
|
||||
"seed helper must create an OpenAI model")
|
||||
|
||||
parent, err := server.CreateChat(ctx, CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "parent-openai",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
parentChat, err := db.GetChatByID(ctx, parent.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
tools := server.subagentTools(ctx, func() database.Chat { return parentChat })
|
||||
tool := findToolByName(tools, "spawn_computer_use_agent")
|
||||
require.NotNil(t, tool)
|
||||
|
||||
resp, err := tool.Run(ctx, fantasy.ToolCall{
|
||||
ID: "call-3",
|
||||
Name: "spawn_computer_use_agent",
|
||||
Input: `{"prompt":"take a screenshot"}`,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.False(t, resp.IsError, "expected success but got: %s", resp.Content)
|
||||
|
||||
// Parse the response to get the child chat ID.
|
||||
var result map[string]any
|
||||
require.NoError(t, json.Unmarshal([]byte(resp.Content), &result))
|
||||
childIDStr, ok := result["chat_id"].(string)
|
||||
require.True(t, ok, "response must contain chat_id")
|
||||
|
||||
childID, err := uuid.Parse(childIDStr)
|
||||
require.NoError(t, err)
|
||||
|
||||
childChat, err := db.GetChatByID(ctx, childID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// The child must have Mode=computer_use which causes
|
||||
// runChat to override the model to the predefined computer
|
||||
// use model instead of using the parent's model config.
|
||||
require.True(t, childChat.Mode.Valid)
|
||||
assert.Equal(t, database.ChatModeComputerUse, childChat.Mode.ChatMode)
|
||||
|
||||
// The predefined computer use model is Anthropic, which
|
||||
// differs from the parent's OpenAI model. This confirms
|
||||
// that the child will not inherit the parent's model at
|
||||
// runtime.
|
||||
assert.NotEqual(t, model.Provider, chattool.ComputerUseModelProvider,
|
||||
"computer use model provider must differ from parent model provider")
|
||||
assert.Equal(t, "anthropic", chattool.ComputerUseModelProvider)
|
||||
assert.NotEmpty(t, chattool.ComputerUseModelName)
|
||||
}
|
||||
@@ -0,0 +1,218 @@
|
||||
package chatd_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/chatd"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/database/dbtestutil"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
)
|
||||
|
||||
func TestSpawnComputerUseAgent_CreatesChildWithChatMode(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
server := newTestServer(t, db, ps, uuid.New())
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
user, model := seedChatDependencies(ctx, t, db)
|
||||
|
||||
// Create a parent chat.
|
||||
parent, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "parent",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Simulate what spawn_computer_use_agent does: set ChatMode
|
||||
// to computer_use and provide a system prompt.
|
||||
prompt := "Use the desktop to open Firefox"
|
||||
|
||||
child, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: parent.OwnerID,
|
||||
ParentChatID: uuid.NullUUID{
|
||||
UUID: parent.ID,
|
||||
Valid: true,
|
||||
},
|
||||
RootChatID: uuid.NullUUID{
|
||||
UUID: parent.ID,
|
||||
Valid: true,
|
||||
},
|
||||
ModelConfigID: model.ID,
|
||||
Title: "computer-use",
|
||||
ChatMode: database.NullChatMode{ChatMode: database.ChatModeComputerUse, Valid: true},
|
||||
SystemPrompt: "Computer use instructions\n\n" + prompt,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText(prompt)},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify parent-child relationship.
|
||||
require.True(t, child.ParentChatID.Valid)
|
||||
require.Equal(t, parent.ID, child.ParentChatID.UUID)
|
||||
|
||||
// Verify the chat type is set correctly.
|
||||
require.True(t, child.Mode.Valid)
|
||||
assert.Equal(t, database.ChatModeComputerUse, child.Mode.ChatMode)
|
||||
|
||||
// Confirm via a fresh DB read as well.
|
||||
got, err := db.GetChatByID(ctx, child.ID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, got.Mode.Valid)
|
||||
assert.Equal(t, database.ChatModeComputerUse, got.Mode.ChatMode)
|
||||
}
|
||||
|
||||
func TestSpawnComputerUseAgent_SystemPromptFormat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
server := newTestServer(t, db, ps, uuid.New())
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
user, model := seedChatDependencies(ctx, t, db)
|
||||
|
||||
parent, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "parent",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
prompt := "Navigate to settings page"
|
||||
systemPrompt := "Computer use instructions\n\n" + prompt
|
||||
|
||||
child, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: parent.OwnerID,
|
||||
ParentChatID: uuid.NullUUID{
|
||||
UUID: parent.ID,
|
||||
Valid: true,
|
||||
},
|
||||
RootChatID: uuid.NullUUID{
|
||||
UUID: parent.ID,
|
||||
Valid: true,
|
||||
},
|
||||
ModelConfigID: model.ID,
|
||||
Title: "computer-use-format",
|
||||
ChatMode: database.NullChatMode{ChatMode: database.ChatModeComputerUse, Valid: true},
|
||||
SystemPrompt: systemPrompt,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText(prompt)},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
messages, err := db.GetChatMessagesForPromptByChatID(ctx, child.ID)
|
||||
require.NoError(t, err)
|
||||
|
||||
// The system message raw content is a JSON-encoded string.
|
||||
// It should contain the system prompt with the user prompt.
|
||||
var rawSystemContent string
|
||||
for _, msg := range messages {
|
||||
if msg.Role != "system" {
|
||||
continue
|
||||
}
|
||||
if msg.Content.Valid {
|
||||
rawSystemContent = string(msg.Content.RawMessage)
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
assert.Contains(t, rawSystemContent, prompt,
|
||||
"system prompt raw content should contain the user prompt")
|
||||
}
|
||||
|
||||
func TestSpawnComputerUseAgent_ChildIsListedUnderParent(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
server := newTestServer(t, db, ps, uuid.New())
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
user, model := seedChatDependencies(ctx, t, db)
|
||||
|
||||
parent, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "parent",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
prompt := "Check the UI layout"
|
||||
|
||||
child, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: parent.OwnerID,
|
||||
ParentChatID: uuid.NullUUID{
|
||||
UUID: parent.ID,
|
||||
Valid: true,
|
||||
},
|
||||
RootChatID: uuid.NullUUID{
|
||||
UUID: parent.ID,
|
||||
Valid: true,
|
||||
},
|
||||
ModelConfigID: model.ID,
|
||||
Title: "computer-use-child",
|
||||
ChatMode: database.NullChatMode{ChatMode: database.ChatModeComputerUse, Valid: true},
|
||||
SystemPrompt: "Computer use instructions\n\n" + prompt,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText(prompt)},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Verify the child is linked to the parent.
|
||||
fetchedChild, err := db.GetChatByID(ctx, child.ID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, fetchedChild.ParentChatID.Valid)
|
||||
assert.Equal(t, parent.ID, fetchedChild.ParentChatID.UUID)
|
||||
}
|
||||
|
||||
func TestSpawnComputerUseAgent_RootChatIDPropagation(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
server := newTestServer(t, db, ps, uuid.New())
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
user, model := seedChatDependencies(ctx, t, db)
|
||||
|
||||
// Create a root parent chat (no parent of its own).
|
||||
parent, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: user.ID,
|
||||
Title: "root-parent",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText("hello")},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
prompt := "Take a screenshot"
|
||||
|
||||
child, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OwnerID: parent.OwnerID,
|
||||
ParentChatID: uuid.NullUUID{
|
||||
UUID: parent.ID,
|
||||
Valid: true,
|
||||
},
|
||||
RootChatID: uuid.NullUUID{
|
||||
UUID: parent.ID,
|
||||
Valid: true,
|
||||
},
|
||||
ModelConfigID: model.ID,
|
||||
Title: "computer-use-root-test",
|
||||
ChatMode: database.NullChatMode{ChatMode: database.ChatModeComputerUse, Valid: true},
|
||||
SystemPrompt: "Computer use instructions\n\n" + prompt,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{codersdk.ChatMessageText(prompt)},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// When the parent has no RootChatID, the child's RootChatID
|
||||
// should point to the parent.
|
||||
require.True(t, child.RootChatID.Valid)
|
||||
assert.Equal(t, parent.ID, child.RootChatID.UUID)
|
||||
|
||||
// Verify chat was retrieved correctly from the DB.
|
||||
got, err := db.GetChatByID(ctx, child.ID)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, got.RootChatID.Valid)
|
||||
assert.Equal(t, parent.ID, got.RootChatID.UUID)
|
||||
}
|
||||
+117
@@ -25,6 +25,7 @@ import (
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/agent/agentssh"
|
||||
"github.com/coder/coder/v2/coderd/audit"
|
||||
"github.com/coder/coder/v2/coderd/chatd"
|
||||
"github.com/coder/coder/v2/coderd/chatd/chatprovider"
|
||||
@@ -43,6 +44,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/searchquery"
|
||||
"github.com/coder/coder/v2/coderd/tracing"
|
||||
"github.com/coder/coder/v2/coderd/util/ptr"
|
||||
"github.com/coder/coder/v2/coderd/workspaceapps"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/codersdk/wsjson"
|
||||
"github.com/coder/websocket"
|
||||
@@ -735,6 +737,121 @@ proxyLoop:
|
||||
_ = clientStream.Close(websocket.StatusGoingAway)
|
||||
}
|
||||
|
||||
// @Summary Watch chat desktop
|
||||
// @ID watch-chat-desktop
|
||||
// @Security CoderSessionToken
|
||||
// @Tags Chats
|
||||
// @Param chat path string true "Chat ID" format(uuid)
|
||||
// @Success 101
|
||||
// @Router /chats/{chat}/desktop [get]
|
||||
//
|
||||
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
|
||||
//
|
||||
//nolint:revive // HTTP handler writes to ResponseWriter.
|
||||
func (api *API) watchChatDesktop(rw http.ResponseWriter, r *http.Request) {
|
||||
var (
|
||||
ctx = r.Context()
|
||||
chat = httpmw.ChatParam(r)
|
||||
logger = api.Logger.Named("chat_desktop").With(slog.F("chat_id", chat.ID))
|
||||
)
|
||||
|
||||
if !chat.WorkspaceID.Valid {
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: "Chat has no workspace.",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
agents, err := api.Database.GetWorkspaceAgentsInLatestBuildByWorkspaceID(ctx, chat.WorkspaceID.UUID)
|
||||
if err != nil {
|
||||
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
|
||||
Message: "Internal error fetching workspace agents.",
|
||||
Detail: err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
if len(agents) == 0 {
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: "Chat workspace has no agents.",
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
apiAgent, err := db2sdk.WorkspaceAgent(
|
||||
api.DERPMap(),
|
||||
*api.TailnetCoordinator.Load(),
|
||||
agents[0],
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
api.AgentInactiveDisconnectTimeout,
|
||||
api.DeploymentValues.AgentFallbackTroubleshootingURL.String(),
|
||||
)
|
||||
if err != nil {
|
||||
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
|
||||
Message: "Internal error reading workspace agent.",
|
||||
Detail: err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
if apiAgent.Status != codersdk.WorkspaceAgentConnected {
|
||||
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
|
||||
Message: fmt.Sprintf("Agent state is %q, must be connected.", apiAgent.Status),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
dialCtx, dialCancel := context.WithTimeout(ctx, 30*time.Second)
|
||||
defer dialCancel()
|
||||
|
||||
agentConn, release, err := api.agentProvider.AgentConn(dialCtx, agents[0].ID)
|
||||
if err != nil {
|
||||
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
|
||||
Message: "Failed to dial workspace agent.",
|
||||
Detail: err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
defer release()
|
||||
|
||||
desktopConn, err := agentConn.ConnectDesktopVNC(ctx)
|
||||
if err != nil {
|
||||
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
|
||||
Message: "Failed to connect to agent desktop.",
|
||||
Detail: err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
defer desktopConn.Close()
|
||||
|
||||
conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
|
||||
CompressionMode: websocket.CompressionDisabled,
|
||||
})
|
||||
if err != nil {
|
||||
logger.Error(ctx, "failed to accept websocket", slog.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
// No read limit — RFB framebuffer updates can be large.
|
||||
conn.SetReadLimit(-1)
|
||||
|
||||
ctx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
ctx, wsNetConn := workspaceapps.WebsocketNetConn(ctx, conn, websocket.MessageBinary)
|
||||
defer wsNetConn.Close()
|
||||
|
||||
go httpapi.HeartbeatClose(ctx, logger, cancel, conn)
|
||||
|
||||
agentssh.Bicopy(ctx, wsNetConn, desktopConn)
|
||||
logger.Debug(ctx, "desktop Bicopy finished")
|
||||
}
|
||||
|
||||
// @Summary Archive a chat
|
||||
// @ID archive-chat
|
||||
// @Tags Chats
|
||||
// @Success 204
|
||||
// @Router /chats/{chat}/archive [post]
|
||||
func (api *API) archiveChat(rw http.ResponseWriter, r *http.Request) {
|
||||
ctx := r.Context()
|
||||
chat := httpmw.ChatParam(r)
|
||||
|
||||
@@ -3802,6 +3802,41 @@ func decRef(value string) *decimal.Decimal {
|
||||
return &d
|
||||
}
|
||||
|
||||
func TestWatchChatDesktop(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("NoWorkspace", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
client := newChatClient(t)
|
||||
_ = coderdtest.CreateFirstUser(t, client)
|
||||
_ = createChatModelConfig(t, client)
|
||||
|
||||
createdChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{
|
||||
Content: []codersdk.ChatInputPart{
|
||||
{
|
||||
Type: codersdk.ChatInputPartTypeText,
|
||||
Text: "desktop no workspace test",
|
||||
},
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Try to connect to the desktop endpoint — should fail because
|
||||
// chat has no workspace.
|
||||
res, err := client.Request(
|
||||
ctx,
|
||||
http.MethodGet,
|
||||
fmt.Sprintf("/api/experimental/chats/%s/desktop", createdChat.ID),
|
||||
nil,
|
||||
)
|
||||
require.NoError(t, err)
|
||||
defer res.Body.Close()
|
||||
require.Equal(t, http.StatusBadRequest, res.StatusCode)
|
||||
})
|
||||
}
|
||||
|
||||
func createChatModelConfig(t *testing.T, client *codersdk.Client) codersdk.ChatModelConfig {
|
||||
t.Helper()
|
||||
|
||||
|
||||
@@ -1179,6 +1179,7 @@ func New(options *Options) *API {
|
||||
r.Use(httpmw.ExtractChatParam(options.Database))
|
||||
r.Get("/", api.getChat)
|
||||
r.Get("/git/watch", api.watchChatGit)
|
||||
r.Get("/desktop", api.watchChatDesktop)
|
||||
r.Post("/archive", api.archiveChat)
|
||||
r.Post("/unarchive", api.unarchiveChat)
|
||||
r.Get("/messages", api.getChatMessages)
|
||||
|
||||
Generated
+6
-1
@@ -278,6 +278,10 @@ CREATE TYPE chat_message_visibility AS ENUM (
|
||||
'both'
|
||||
);
|
||||
|
||||
CREATE TYPE chat_mode AS ENUM (
|
||||
'computer_use'
|
||||
);
|
||||
|
||||
CREATE TYPE chat_status AS ENUM (
|
||||
'waiting',
|
||||
'pending',
|
||||
@@ -1306,7 +1310,8 @@ CREATE TABLE chats (
|
||||
root_chat_id uuid,
|
||||
last_model_config_id uuid NOT NULL,
|
||||
archived boolean DEFAULT false NOT NULL,
|
||||
last_error text
|
||||
last_error text,
|
||||
mode chat_mode
|
||||
);
|
||||
|
||||
CREATE TABLE connection_logs (
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
ALTER TABLE chats DROP COLUMN mode;
|
||||
DROP TYPE IF EXISTS chat_mode;
|
||||
@@ -0,0 +1,3 @@
|
||||
CREATE TYPE chat_mode AS ENUM ('computer_use');
|
||||
|
||||
ALTER TABLE chats ADD COLUMN mode chat_mode;
|
||||
@@ -1174,6 +1174,61 @@ func AllChatMessageVisibilityValues() []ChatMessageVisibility {
|
||||
}
|
||||
}
|
||||
|
||||
type ChatMode string
|
||||
|
||||
const (
|
||||
ChatModeComputerUse ChatMode = "computer_use"
|
||||
)
|
||||
|
||||
func (e *ChatMode) Scan(src interface{}) error {
|
||||
switch s := src.(type) {
|
||||
case []byte:
|
||||
*e = ChatMode(s)
|
||||
case string:
|
||||
*e = ChatMode(s)
|
||||
default:
|
||||
return fmt.Errorf("unsupported scan type for ChatMode: %T", src)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type NullChatMode struct {
|
||||
ChatMode ChatMode `json:"chat_mode"`
|
||||
Valid bool `json:"valid"` // Valid is true if ChatMode is not NULL
|
||||
}
|
||||
|
||||
// Scan implements the Scanner interface.
|
||||
func (ns *NullChatMode) Scan(value interface{}) error {
|
||||
if value == nil {
|
||||
ns.ChatMode, ns.Valid = "", false
|
||||
return nil
|
||||
}
|
||||
ns.Valid = true
|
||||
return ns.ChatMode.Scan(value)
|
||||
}
|
||||
|
||||
// Value implements the driver Valuer interface.
|
||||
func (ns NullChatMode) Value() (driver.Value, error) {
|
||||
if !ns.Valid {
|
||||
return nil, nil
|
||||
}
|
||||
return string(ns.ChatMode), nil
|
||||
}
|
||||
|
||||
func (e ChatMode) Valid() bool {
|
||||
switch e {
|
||||
case ChatModeComputerUse:
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func AllChatModeValues() []ChatMode {
|
||||
return []ChatMode{
|
||||
ChatModeComputerUse,
|
||||
}
|
||||
}
|
||||
|
||||
type ChatStatus string
|
||||
|
||||
const (
|
||||
@@ -3972,6 +4027,7 @@ type Chat struct {
|
||||
LastModelConfigID uuid.UUID `db:"last_model_config_id" json:"last_model_config_id"`
|
||||
Archived bool `db:"archived" json:"archived"`
|
||||
LastError sql.NullString `db:"last_error" json:"last_error"`
|
||||
Mode NullChatMode `db:"mode" json:"mode"`
|
||||
}
|
||||
|
||||
type ChatDiffStatus struct {
|
||||
|
||||
@@ -2949,7 +2949,7 @@ WHERE
|
||||
$3::int
|
||||
)
|
||||
RETURNING
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode
|
||||
`
|
||||
|
||||
type AcquireChatsParams struct {
|
||||
@@ -2985,6 +2985,7 @@ func (q *sqlQuerier) AcquireChats(ctx context.Context, arg AcquireChatsParams) (
|
||||
&i.LastModelConfigID,
|
||||
&i.Archived,
|
||||
&i.LastError,
|
||||
&i.Mode,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -3172,7 +3173,7 @@ func (q *sqlQuerier) DeleteChatQueuedMessage(ctx context.Context, arg DeleteChat
|
||||
|
||||
const getChatByID = `-- name: GetChatByID :one
|
||||
SELECT
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode
|
||||
FROM
|
||||
chats
|
||||
WHERE
|
||||
@@ -3198,12 +3199,13 @@ func (q *sqlQuerier) GetChatByID(ctx context.Context, id uuid.UUID) (Chat, error
|
||||
&i.LastModelConfigID,
|
||||
&i.Archived,
|
||||
&i.LastError,
|
||||
&i.Mode,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getChatByIDForUpdate = `-- name: GetChatByIDForUpdate :one
|
||||
SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error FROM chats WHERE id = $1::uuid FOR UPDATE
|
||||
SELECT id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode FROM chats WHERE id = $1::uuid FOR UPDATE
|
||||
`
|
||||
|
||||
func (q *sqlQuerier) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (Chat, error) {
|
||||
@@ -3225,6 +3227,7 @@ func (q *sqlQuerier) GetChatByIDForUpdate(ctx context.Context, id uuid.UUID) (Ch
|
||||
&i.LastModelConfigID,
|
||||
&i.Archived,
|
||||
&i.LastError,
|
||||
&i.Mode,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -3886,7 +3889,7 @@ func (q *sqlQuerier) GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID
|
||||
|
||||
const getChatsByOwnerID = `-- name: GetChatsByOwnerID :many
|
||||
SELECT
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode
|
||||
FROM
|
||||
chats
|
||||
WHERE
|
||||
@@ -3963,6 +3966,7 @@ func (q *sqlQuerier) GetChatsByOwnerID(ctx context.Context, arg GetChatsByOwnerI
|
||||
&i.LastModelConfigID,
|
||||
&i.Archived,
|
||||
&i.LastError,
|
||||
&i.Mode,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -4024,7 +4028,7 @@ func (q *sqlQuerier) GetLastChatMessageByRole(ctx context.Context, arg GetLastCh
|
||||
|
||||
const getStaleChats = `-- name: GetStaleChats :many
|
||||
SELECT
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode
|
||||
FROM
|
||||
chats
|
||||
WHERE
|
||||
@@ -4059,6 +4063,7 @@ func (q *sqlQuerier) GetStaleChats(ctx context.Context, staleThreshold time.Time
|
||||
&i.LastModelConfigID,
|
||||
&i.Archived,
|
||||
&i.LastError,
|
||||
&i.Mode,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -4080,17 +4085,19 @@ INSERT INTO chats (
|
||||
parent_chat_id,
|
||||
root_chat_id,
|
||||
last_model_config_id,
|
||||
title
|
||||
title,
|
||||
mode
|
||||
) VALUES (
|
||||
$1::uuid,
|
||||
$2::uuid,
|
||||
$3::uuid,
|
||||
$4::uuid,
|
||||
$5::uuid,
|
||||
$6::text
|
||||
$6::text,
|
||||
$7::chat_mode
|
||||
)
|
||||
RETURNING
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode
|
||||
`
|
||||
|
||||
type InsertChatParams struct {
|
||||
@@ -4100,6 +4107,7 @@ type InsertChatParams struct {
|
||||
RootChatID uuid.NullUUID `db:"root_chat_id" json:"root_chat_id"`
|
||||
LastModelConfigID uuid.UUID `db:"last_model_config_id" json:"last_model_config_id"`
|
||||
Title string `db:"title" json:"title"`
|
||||
Mode NullChatMode `db:"mode" json:"mode"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) InsertChat(ctx context.Context, arg InsertChatParams) (Chat, error) {
|
||||
@@ -4110,6 +4118,7 @@ func (q *sqlQuerier) InsertChat(ctx context.Context, arg InsertChatParams) (Chat
|
||||
arg.RootChatID,
|
||||
arg.LastModelConfigID,
|
||||
arg.Title,
|
||||
arg.Mode,
|
||||
)
|
||||
var i Chat
|
||||
err := row.Scan(
|
||||
@@ -4128,6 +4137,7 @@ func (q *sqlQuerier) InsertChat(ctx context.Context, arg InsertChatParams) (Chat
|
||||
&i.LastModelConfigID,
|
||||
&i.Archived,
|
||||
&i.LastError,
|
||||
&i.Mode,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -4307,7 +4317,7 @@ SET
|
||||
WHERE
|
||||
id = $2::uuid
|
||||
RETURNING
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode
|
||||
`
|
||||
|
||||
type UpdateChatByIDParams struct {
|
||||
@@ -4334,6 +4344,7 @@ func (q *sqlQuerier) UpdateChatByID(ctx context.Context, arg UpdateChatByIDParam
|
||||
&i.LastModelConfigID,
|
||||
&i.Archived,
|
||||
&i.LastError,
|
||||
&i.Mode,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -4421,7 +4432,7 @@ SET
|
||||
WHERE
|
||||
id = $6::uuid
|
||||
RETURNING
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode
|
||||
`
|
||||
|
||||
type UpdateChatStatusParams struct {
|
||||
@@ -4459,6 +4470,7 @@ func (q *sqlQuerier) UpdateChatStatus(ctx context.Context, arg UpdateChatStatusP
|
||||
&i.LastModelConfigID,
|
||||
&i.Archived,
|
||||
&i.LastError,
|
||||
&i.Mode,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
@@ -4472,7 +4484,7 @@ SET
|
||||
WHERE
|
||||
id = $2::uuid
|
||||
RETURNING
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error
|
||||
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode
|
||||
`
|
||||
|
||||
type UpdateChatWorkspaceParams struct {
|
||||
@@ -4499,6 +4511,7 @@ func (q *sqlQuerier) UpdateChatWorkspace(ctx context.Context, arg UpdateChatWork
|
||||
&i.LastModelConfigID,
|
||||
&i.Archived,
|
||||
&i.LastError,
|
||||
&i.Mode,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
@@ -142,14 +142,16 @@ INSERT INTO chats (
|
||||
parent_chat_id,
|
||||
root_chat_id,
|
||||
last_model_config_id,
|
||||
title
|
||||
title,
|
||||
mode
|
||||
) VALUES (
|
||||
@owner_id::uuid,
|
||||
sqlc.narg('workspace_id')::uuid,
|
||||
sqlc.narg('parent_chat_id')::uuid,
|
||||
sqlc.narg('root_chat_id')::uuid,
|
||||
@last_model_config_id::uuid,
|
||||
@title::text
|
||||
@title::text,
|
||||
sqlc.narg('mode')::chat_mode
|
||||
)
|
||||
RETURNING
|
||||
*;
|
||||
|
||||
Reference in New Issue
Block a user