Skip to content

Commit 69d1a54

Browse files
DanielleMaywoodDanielleMaywood
committed
refactor: modify task creation endpoint to return a task, not workspace
1 parent f721f3d commit 69d1a54

File tree

11 files changed

+414
-111
lines changed

11 files changed

+414
-111
lines changed

‎cli/exp_taskcreate.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,7 +104,7 @@ func (r *RootCmd) taskCreate() *serpent.Command {
104104
templateVersionPresetID = preset.ID
105105
}
106106

107-
workspace, err := expClient.CreateTask(ctx, codersdk.Me, codersdk.CreateTaskRequest{
107+
task, err := expClient.CreateTask(ctx, codersdk.Me, codersdk.CreateTaskRequest{
108108
TemplateVersionID: templateVersionID,
109109
TemplateVersionPresetID: templateVersionPresetID,
110110
Prompt: taskInput,
@@ -116,8 +116,8 @@ func (r *RootCmd) taskCreate() *serpent.Command {
116116
_, _ = fmt.Fprintf(
117117
inv.Stdout,
118118
"The task %s has been created at %s!\n",
119-
cliui.Keyword(workspace.Name),
120-
cliui.Timestamp(workspace.CreatedAt),
119+
cliui.Keyword(task.Name),
120+
cliui.Timestamp(task.CreatedAt),
121121
)
122122

123123
return nil

‎coderd/aitasks.go

Lines changed: 82 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package coderd
33
import (
44
"context"
55
"database/sql"
6+
"encoding/json"
67
"errors"
78
"fmt"
89
"net/http"
@@ -22,6 +23,7 @@ import (
2223
"github.com/coder/coder/v2/coderd/rbac/policy"
2324
"github.com/coder/coder/v2/coderd/searchquery"
2425
"github.com/coder/coder/v2/coderd/taskname"
26+
"github.com/coder/coder/v2/coderd/util/rwsink"
2527
"github.com/coder/coder/v2/codersdk"
2628
)
2729

@@ -186,9 +188,86 @@ func (api *API) tasksCreate(rw http.ResponseWriter, r *http.Request) {
186188
WorkspaceOwner: owner.Username,
187189
},
188190
})
189-
190191
defer commitAudit()
191-
createWorkspace(ctx, aReq, apiKey.UserID, api, owner, createReq, rw, r)
192+
193+
rwSink := rwsink.New()
194+
createWorkspace(ctx, aReq, apiKey.UserID, api, owner, createReq, rwSink, r)
195+
196+
if rwSink.StatusCode != nil && *rwSink.StatusCode == http.StatusCreated {
197+
bytes := rwSink.ResetBody()
198+
199+
var ws codersdk.Workspace
200+
if err := json.Unmarshal(bytes, &ws); err != nil {
201+
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
202+
Message: "Internal error decoding created workspace",
203+
Detail: err.Error(),
204+
})
205+
return
206+
}
207+
208+
task := taskFromWorkspace(ws, req.Prompt)
209+
210+
httpapi.Write(ctx, rwSink, http.StatusCreated, task)
211+
}
212+
213+
_ = rwSink.FlushTo(rw)
214+
}
215+
216+
func taskFromWorkspace(ws codersdk.Workspace, initialPrompt string) codersdk.Task {
217+
// TODO(DanielleMaywood):
218+
// This just picks up the first agent it discovers.
219+
// This approach _might_ break when a task has multiple agents,
220+
// depending on which agent was found first.
221+
//
222+
// We explicitly do not have support for running tasks
223+
// inside of a sub agent at the moment, so we can be sure
224+
// that any sub agents are not the agent we're looking for.
225+
var taskAgentID uuid.NullUUID
226+
var taskAgentLifecycle *codersdk.WorkspaceAgentLifecycle
227+
var taskAgentHealth *codersdk.WorkspaceAgentHealth
228+
for _, resource := range ws.LatestBuild.Resources {
229+
for _, agent := range resource.Agents {
230+
if agent.ParentID.Valid {
231+
continue
232+
}
233+
234+
taskAgentID = uuid.NullUUID{Valid: true, UUID: agent.ID}
235+
taskAgentLifecycle = &agent.LifecycleState
236+
taskAgentHealth = &agent.Health
237+
break
238+
}
239+
}
240+
241+
var currentState *codersdk.TaskStateEntry
242+
if ws.LatestAppStatus != nil {
243+
currentState = &codersdk.TaskStateEntry{
244+
Timestamp: ws.LatestAppStatus.CreatedAt,
245+
State: codersdk.TaskState(ws.LatestAppStatus.State),
246+
Message: ws.LatestAppStatus.Message,
247+
URI: ws.LatestAppStatus.URI,
248+
}
249+
}
250+
251+
return codersdk.Task{
252+
ID: ws.ID,
253+
OrganizationID: ws.OrganizationID,
254+
OwnerID: ws.OwnerID,
255+
OwnerName: ws.OwnerName,
256+
Name: ws.Name,
257+
TemplateID: ws.TemplateID,
258+
TemplateName: ws.TemplateName,
259+
TemplateDisplayName: ws.TemplateDisplayName,
260+
TemplateIcon: ws.TemplateIcon,
261+
WorkspaceID: uuid.NullUUID{Valid: true, UUID: ws.ID},
262+
WorkspaceAgentID: taskAgentID,
263+
WorkspaceAgentLifecycle: taskAgentLifecycle,
264+
WorkspaceAgentHealth: taskAgentHealth,
265+
CreatedAt: ws.CreatedAt,
266+
UpdatedAt: ws.UpdatedAt,
267+
InitialPrompt: initialPrompt,
268+
Status: ws.LatestBuild.Status,
269+
CurrentState: currentState,
270+
}
192271
}
193272

194273
// tasksFromWorkspaces converts a slice of API workspaces into tasks, fetching
@@ -213,60 +292,7 @@ func (api *API) tasksFromWorkspaces(ctx context.Context, apiWorkspaces []codersd
213292

214293
tasks := make([]codersdk.Task, 0, len(apiWorkspaces))
215294
for _, ws := range apiWorkspaces {
216-
// TODO(DanielleMaywood):
217-
// This just picks up the first agent it discovers.
218-
// This approach _might_ break when a task has multiple agents,
219-
// depending on which agent was found first.
220-
//
221-
// We explicitly do not have support for running tasks
222-
// inside of a sub agent at the moment, so we can be sure
223-
// that any sub agents are not the agent we're looking for.
224-
var taskAgentID uuid.NullUUID
225-
var taskAgentLifecycle *codersdk.WorkspaceAgentLifecycle
226-
var taskAgentHealth *codersdk.WorkspaceAgentHealth
227-
for _, resource := range ws.LatestBuild.Resources {
228-
for _, agent := range resource.Agents {
229-
if agent.ParentID.Valid {
230-
continue
231-
}
232-
233-
taskAgentID = uuid.NullUUID{Valid: true, UUID: agent.ID}
234-
taskAgentLifecycle = &agent.LifecycleState
235-
taskAgentHealth = &agent.Health
236-
break
237-
}
238-
}
239-
240-
var currentState *codersdk.TaskStateEntry
241-
if ws.LatestAppStatus != nil {
242-
currentState = &codersdk.TaskStateEntry{
243-
Timestamp: ws.LatestAppStatus.CreatedAt,
244-
State: codersdk.TaskState(ws.LatestAppStatus.State),
245-
Message: ws.LatestAppStatus.Message,
246-
URI: ws.LatestAppStatus.URI,
247-
}
248-
}
249-
250-
tasks = append(tasks, codersdk.Task{
251-
ID: ws.ID,
252-
OrganizationID: ws.OrganizationID,
253-
OwnerID: ws.OwnerID,
254-
OwnerName: ws.OwnerName,
255-
Name: ws.Name,
256-
TemplateID: ws.TemplateID,
257-
TemplateName: ws.TemplateName,
258-
TemplateDisplayName: ws.TemplateDisplayName,
259-
TemplateIcon: ws.TemplateIcon,
260-
WorkspaceID: uuid.NullUUID{Valid: true, UUID: ws.ID},
261-
WorkspaceAgentID: taskAgentID,
262-
WorkspaceAgentLifecycle: taskAgentLifecycle,
263-
WorkspaceAgentHealth: taskAgentHealth,
264-
CreatedAt: ws.CreatedAt,
265-
UpdatedAt: ws.UpdatedAt,
266-
InitialPrompt: promptsByBuildID[ws.LatestBuild.ID],
267-
Status: ws.LatestBuild.Status,
268-
CurrentState: currentState,
269-
})
295+
tasks = append(tasks, taskFromWorkspace(ws, promptsByBuildID[ws.LatestBuild.ID]))
270296
}
271297

272298
return tasks, nil

‎coderd/aitasks_test.go

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -299,19 +299,23 @@ func TestTasksCreate(t *testing.T) {
299299
expClient := codersdk.NewExperimentalClient(client)
300300

301301
// When: We attempt to create a Task.
302-
workspace, err := expClient.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
302+
task, err := expClient.CreateTask(ctx, "me", codersdk.CreateTaskRequest{
303303
TemplateVersionID: template.ActiveVersionID,
304304
Prompt: taskPrompt,
305305
})
306306
require.NoError(t, err)
307-
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID)
307+
require.True(t, task.WorkspaceID.Valid)
308+
309+
ws, err := client.Workspace(ctx, task.WorkspaceID.UUID)
310+
require.NoError(t, err)
311+
coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, ws.LatestBuild.ID)
308312

309313
// Then: We expect a workspace to have been created.
310-
assert.NotEmpty(t, workspace.Name)
311-
assert.Equal(t, template.ID, workspace.TemplateID)
314+
assert.NotEmpty(t, task.Name)
315+
assert.Equal(t, template.ID, task.TemplateID)
312316

313317
// And: We expect it to have the "AI Prompt" parameter correctly set.
314-
parameters, err := client.WorkspaceBuildParameters(ctx, workspace.LatestBuild.ID)
318+
parameters, err := client.WorkspaceBuildParameters(ctx, ws.LatestBuild.ID)
315319
require.NoError(t, err)
316320
require.Len(t, parameters, 1)
317321
assert.Equal(t, codersdk.AITaskPromptParameterName, parameters[0].Name)

‎coderd/util/rwsink/rwsink.go

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
package rwsink
2+
3+
import (
4+
"bytes"
5+
"maps"
6+
"net/http"
7+
"slices"
8+
)
9+
10+
type ResponseWriterSink struct {
11+
StatusCode *int
12+
header http.Header
13+
body bytes.Buffer
14+
}
15+
16+
var _ http.ResponseWriter = (*ResponseWriterSink)(nil)
17+
18+
func New() *ResponseWriterSink {
19+
return &ResponseWriterSink{
20+
StatusCode: nil,
21+
header: make(http.Header),
22+
body: bytes.Buffer{},
23+
}
24+
}
25+
26+
func (rw *ResponseWriterSink) Header() http.Header {
27+
return rw.header
28+
}
29+
30+
func (rw *ResponseWriterSink) Write(data []byte) (int, error) {
31+
return rw.body.Write(data)
32+
}
33+
34+
func (rw *ResponseWriterSink) ResetBody() []byte {
35+
body := slices.Clone(rw.body.Bytes())
36+
rw.body = bytes.Buffer{}
37+
return body
38+
}
39+
40+
func (rw *ResponseWriterSink) WriteHeader(statusCode int) {
41+
if rw.StatusCode == nil {
42+
rw.StatusCode = &statusCode
43+
}
44+
}
45+
46+
func (rw *ResponseWriterSink) Bytes() []byte {
47+
return rw.body.Bytes()
48+
}
49+
50+
func (rw *ResponseWriterSink) FlushTo(to http.ResponseWriter) error {
51+
maps.Copy(to.Header(), rw.header)
52+
53+
if rw.StatusCode != nil {
54+
to.WriteHeader(*rw.StatusCode)
55+
}
56+
57+
if _, err := to.Write(rw.body.Bytes()); err != nil {
58+
return err
59+
}
60+
61+
return nil
62+
}

0 commit comments

Comments
 (0)