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:
Hugo Dutka
2026-03-13 19:49:34 +01:00
committed by GitHub
parent 67f5494665
commit 84527390c6
34 changed files with 4465 additions and 49 deletions
+43
View File
@@ -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": [
+39
View File
@@ -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
View File
@@ -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",
},
})
}
+354
View File
@@ -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)
}
+33 -11
View File
@@ -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
}
+220
View File
@@ -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)
})
}
}
+186
View File
@@ -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
View File
@@ -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) {
+300
View File
@@ -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)
}
+218
View File
@@ -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
View File
@@ -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)
+35
View File
@@ -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()
+1
View File
@@ -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)
+6 -1
View File
@@ -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;
+56
View File
@@ -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 {
+24 -11
View File
@@ -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
}
+4 -2
View File
@@ -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
*;