Skip to content

Commit 896ae36

Browse files
hugodutkahugodutka
committed
endpoint
1 parent 73d50c4 commit 896ae36

File tree

4 files changed

+254
-0
lines changed

4 files changed

+254
-0
lines changed

‎coderd/aitasks.go

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
package coderd
2+
3+
import (
4+
"fmt"
5+
"net/http"
6+
"strings"
7+
8+
"github.com/google/uuid"
9+
10+
"github.com/coder/coder/v2/coderd/httpapi"
11+
"github.com/coder/coder/v2/codersdk"
12+
)
13+
14+
// @Summary Get AI tasks prompts
15+
// @ID get-ai-tasks-prompts
16+
// @Security CoderSessionToken
17+
// @Produce json
18+
// @Tags AITasks
19+
// @Param build_ids query string true "Comma-separated workspace build IDs" format(uuid)
20+
// @Success 200 {object} codersdk.AITasksPromptsResponse
21+
// @Router /aitasks/prompts [get]
22+
func (api *API) aiTasksPrompts(rw http.ResponseWriter, r *http.Request) {
23+
ctx := r.Context()
24+
25+
buildIDsParam := r.URL.Query().Get("build_ids")
26+
if buildIDsParam == "" {
27+
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
28+
Message: "build_ids query parameter is required",
29+
})
30+
return
31+
}
32+
33+
// Parse build IDs
34+
buildIDStrings := strings.Split(buildIDsParam, ",")
35+
buildIDs := make([]uuid.UUID, 0, len(buildIDStrings))
36+
for _, idStr := range buildIDStrings {
37+
id, err := uuid.Parse(strings.TrimSpace(idStr))
38+
if err != nil {
39+
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
40+
Message: fmt.Sprintf("Invalid build ID format: %s", idStr),
41+
Detail: err.Error(),
42+
})
43+
return
44+
}
45+
buildIDs = append(buildIDs, id)
46+
}
47+
48+
parameters, err := api.Database.GetWorkspaceBuildParametersByBuildIDs(ctx, buildIDs)
49+
if err != nil {
50+
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
51+
Message: "Internal error fetching workspace build parameters.",
52+
Detail: err.Error(),
53+
})
54+
return
55+
}
56+
57+
promptsByBuildID := make(map[string]string, len(parameters))
58+
for _, param := range parameters {
59+
if param.Name != codersdk.AITaskPromptParameterName {
60+
continue
61+
}
62+
buildID := param.WorkspaceBuildID.String()
63+
promptsByBuildID[buildID] = param.Value
64+
}
65+
66+
httpapi.Write(ctx, rw, http.StatusOK, codersdk.AITasksPromptsResponse{
67+
Prompts: promptsByBuildID,
68+
})
69+
}

‎coderd/aitasks_test.go

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
package coderd_test
2+
3+
import (
4+
"testing"
5+
6+
"github.com/google/uuid"
7+
"github.com/stretchr/testify/require"
8+
9+
"github.com/coder/coder/v2/coderd/coderdtest"
10+
"github.com/coder/coder/v2/coderd/database/dbtestutil"
11+
"github.com/coder/coder/v2/codersdk"
12+
"github.com/coder/coder/v2/provisioner/echo"
13+
"github.com/coder/coder/v2/provisionersdk/proto"
14+
"github.com/coder/coder/v2/testutil"
15+
)
16+
17+
func TestAITasksPrompts(t *testing.T) {
18+
t.Parallel()
19+
20+
t.Run("EmptyBuildIDs", func(t *testing.T) {
21+
t.Parallel()
22+
client := coderdtest.New(t, &coderdtest.Options{})
23+
_ = coderdtest.CreateFirstUser(t, client)
24+
25+
ctx := testutil.Context(t, testutil.WaitShort)
26+
27+
// Test with empty build IDs
28+
prompts, err := client.AITaskPrompts(ctx, []uuid.UUID{})
29+
require.NoError(t, err)
30+
require.Empty(t, prompts.Prompts)
31+
})
32+
33+
t.Run("MultipleBuilds", func(t *testing.T) {
34+
t.Parallel()
35+
36+
if !dbtestutil.WillUsePostgres() {
37+
t.Skip("This test checks RBAC, which is not supported in the in-memory database")
38+
}
39+
40+
adminClient := coderdtest.New(t, &coderdtest.Options{IncludeProvisionerDaemon: true})
41+
first := coderdtest.CreateFirstUser(t, adminClient)
42+
memberClient, _ := coderdtest.CreateAnotherUser(t, adminClient, first.OrganizationID)
43+
44+
ctx := testutil.Context(t, testutil.WaitLong)
45+
46+
// Create a template with parameters
47+
version := coderdtest.CreateTemplateVersion(t, adminClient, first.OrganizationID, &echo.Responses{
48+
Parse: echo.ParseComplete,
49+
ProvisionPlan: []*proto.Response{{
50+
Type: &proto.Response_Plan{
51+
Plan: &proto.PlanComplete{
52+
Parameters: []*proto.RichParameter{
53+
{
54+
Name: "param1",
55+
Type: "string",
56+
DefaultValue: "default1",
57+
},
58+
{
59+
Name: codersdk.AITaskPromptParameterName,
60+
Type: "string",
61+
DefaultValue: "default2",
62+
},
63+
},
64+
},
65+
},
66+
}},
67+
ProvisionApply: echo.ApplyComplete,
68+
})
69+
template := coderdtest.CreateTemplate(t, adminClient, first.OrganizationID, version.ID)
70+
coderdtest.AwaitTemplateVersionJobCompleted(t, adminClient, version.ID)
71+
72+
// Create two workspaces with different parameters
73+
workspace1 := coderdtest.CreateWorkspace(t, memberClient, template.ID, func(request *codersdk.CreateWorkspaceRequest) {
74+
request.RichParameterValues = []codersdk.WorkspaceBuildParameter{
75+
{Name: "param1", Value: "value1a"},
76+
{Name: codersdk.AITaskPromptParameterName, Value: "value2a"},
77+
}
78+
})
79+
coderdtest.AwaitWorkspaceBuildJobCompleted(t, memberClient, workspace1.LatestBuild.ID)
80+
81+
workspace2 := coderdtest.CreateWorkspace(t, memberClient, template.ID, func(request *codersdk.CreateWorkspaceRequest) {
82+
request.RichParameterValues = []codersdk.WorkspaceBuildParameter{
83+
{Name: "param1", Value: "value1b"},
84+
{Name: codersdk.AITaskPromptParameterName, Value: "value2b"},
85+
}
86+
})
87+
coderdtest.AwaitWorkspaceBuildJobCompleted(t, memberClient, workspace2.LatestBuild.ID)
88+
89+
workspace3 := coderdtest.CreateWorkspace(t, adminClient, template.ID, func(request *codersdk.CreateWorkspaceRequest) {
90+
request.RichParameterValues = []codersdk.WorkspaceBuildParameter{
91+
{Name: "param1", Value: "value1c"},
92+
{Name: codersdk.AITaskPromptParameterName, Value: "value2c"},
93+
}
94+
})
95+
coderdtest.AwaitWorkspaceBuildJobCompleted(t, adminClient, workspace3.LatestBuild.ID)
96+
allBuildIDs := []uuid.UUID{workspace1.LatestBuild.ID, workspace2.LatestBuild.ID, workspace3.LatestBuild.ID}
97+
98+
// Test parameters endpoint as member
99+
prompts, err := memberClient.AITaskPrompts(ctx, allBuildIDs)
100+
require.NoError(t, err)
101+
// we expect 2 prompts because the member client does not have access to workspace3
102+
// since it was created by the admin client
103+
require.Len(t, prompts.Prompts, 2)
104+
105+
// Check workspace1 parameters
106+
build1Prompt := prompts.Prompts[workspace1.LatestBuild.ID.String()]
107+
require.Equal(t, "value2a", build1Prompt)
108+
109+
// Check workspace2 parameters
110+
build2Prompt := prompts.Prompts[workspace2.LatestBuild.ID.String()]
111+
require.Equal(t, "value2b", build2Prompt)
112+
113+
// Test parameters endpoint as admin
114+
// we expect 3 prompts because the admin client has access to all workspaces
115+
prompts, err = adminClient.AITaskPrompts(ctx, allBuildIDs)
116+
require.NoError(t, err)
117+
require.Len(t, prompts.Prompts, 3)
118+
119+
// Check workspace3 parameters
120+
build3Prompt := prompts.Prompts[workspace3.LatestBuild.ID.String()]
121+
require.Equal(t, "value2c", build3Prompt)
122+
})
123+
124+
t.Run("NonExistentBuildIDs", func(t *testing.T) {
125+
t.Parallel()
126+
client := coderdtest.New(t, &coderdtest.Options{})
127+
_ = coderdtest.CreateFirstUser(t, client)
128+
129+
ctx := testutil.Context(t, testutil.WaitShort)
130+
131+
// Test with non-existent build IDs
132+
nonExistentID := uuid.New()
133+
prompts, err := client.AITaskPrompts(ctx, []uuid.UUID{nonExistentID})
134+
require.NoError(t, err)
135+
require.Empty(t, prompts.Prompts)
136+
})
137+
}

‎coderd/coderd.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1494,6 +1494,10 @@ func New(options *Options) *API {
14941494
r.Use(apiKeyMiddleware)
14951495
r.Get("/", api.tailnetRPCConn)
14961496
})
1497+
r.Route("/aitasks", func(r chi.Router) {
1498+
r.Use(apiKeyMiddleware)
1499+
r.Get("/prompts", api.aiTasksPrompts)
1500+
})
14971501
})
14981502

14991503
if options.SwaggerEndpoint {

‎codersdk/aitasks.go

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package codersdk
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"net/http"
7+
"strings"
8+
9+
"github.com/google/uuid"
10+
)
11+
12+
const AITaskPromptParameterName = "AI Prompt"
13+
14+
type AITasksPromptsResponse struct {
15+
// Prompts is a map of workspace build IDs to prompts.
16+
Prompts map[string]string `json:"prompts"`
17+
}
18+
19+
// AITaskPrompts returns prompts for multiple workspace builds by their IDs.
20+
func (c *Client) AITaskPrompts(ctx context.Context, buildIDs []uuid.UUID) (AITasksPromptsResponse, error) {
21+
if len(buildIDs) == 0 {
22+
return AITasksPromptsResponse{
23+
Prompts: make(map[string]string),
24+
}, nil
25+
}
26+
27+
// Convert UUIDs to strings and join them
28+
buildIDStrings := make([]string, len(buildIDs))
29+
for i, id := range buildIDs {
30+
buildIDStrings[i] = id.String()
31+
}
32+
buildIDsParam := strings.Join(buildIDStrings, ",")
33+
34+
res, err := c.Request(ctx, http.MethodGet, "/api/v2/aitasks/prompts", nil, WithQueryParam("build_ids", buildIDsParam))
35+
if err != nil {
36+
return AITasksPromptsResponse{}, err
37+
}
38+
defer res.Body.Close()
39+
if res.StatusCode != http.StatusOK {
40+
return AITasksPromptsResponse{}, ReadBodyAsError(res)
41+
}
42+
var prompts AITasksPromptsResponse
43+
return prompts, json.NewDecoder(res.Body).Decode(&prompts)
44+
}

0 commit comments

Comments
 (0)