Skip to content

Commit b5bd455

Browse files
hugodutkahugodutka
committed
endpoint
1 parent b5260f5 commit b5bd455

File tree

5 files changed

+272
-0
lines changed

5 files changed

+272
-0
lines changed

‎coderd/aitasks.go

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
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+
// This endpoint is experimental and not guaranteed to be stable, so we're not
15+
// generating public-facing documentation for it.
16+
func (api *API) aiTasksPrompts(rw http.ResponseWriter, r *http.Request) {
17+
ctx := r.Context()
18+
19+
buildIDsParam := r.URL.Query().Get("build_ids")
20+
if buildIDsParam == "" {
21+
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
22+
Message: "build_ids query parameter is required",
23+
})
24+
return
25+
}
26+
27+
// Parse build IDs
28+
buildIDStrings := strings.Split(buildIDsParam, ",")
29+
buildIDs := make([]uuid.UUID, 0, len(buildIDStrings))
30+
for _, idStr := range buildIDStrings {
31+
id, err := uuid.Parse(strings.TrimSpace(idStr))
32+
if err != nil {
33+
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
34+
Message: fmt.Sprintf("Invalid build ID format: %s", idStr),
35+
Detail: err.Error(),
36+
})
37+
return
38+
}
39+
buildIDs = append(buildIDs, id)
40+
}
41+
42+
parameters, err := api.Database.GetWorkspaceBuildParametersByBuildIDs(ctx, buildIDs)
43+
if err != nil {
44+
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
45+
Message: "Internal error fetching workspace build parameters.",
46+
Detail: err.Error(),
47+
})
48+
return
49+
}
50+
51+
promptsByBuildID := make(map[string]string, len(parameters))
52+
for _, param := range parameters {
53+
if param.Name != codersdk.AITaskPromptParameterName {
54+
continue
55+
}
56+
buildID := param.WorkspaceBuildID.String()
57+
promptsByBuildID[buildID] = param.Value
58+
}
59+
60+
httpapi.Write(ctx, rw, http.StatusOK, codersdk.AITasksPromptsResponse{
61+
Prompts: promptsByBuildID,
62+
})
63+
}

‎coderd/aitasks_test.go

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

‎coderd/coderd.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -939,6 +939,14 @@ func New(options *Options) *API {
939939
})
940940
})
941941

942+
// Experimental routes are not guaranteed to be stable and may change at any time.
943+
r.Route("/api/experimental", func(r chi.Router) {
944+
r.Use(apiKeyMiddleware)
945+
r.Route("/aitasks", func(r chi.Router) {
946+
r.Get("/prompts", api.aiTasksPrompts)
947+
})
948+
})
949+
942950
r.Route("/api/v2", func(r chi.Router) {
943951
api.APIHandler = r
944952

‎codersdk/aitasks.go

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

‎codersdk/client_experimental.go

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
package codersdk
2+
3+
// ExperimentalClient is a client for the experimental API.
4+
// Its interface is not guaranteed to be stable and may change at any time.
5+
// @typescript-ignore ExperimentalClient
6+
type ExperimentalClient struct {
7+
*Client
8+
}
9+
10+
func NewExperimentalClient(client *Client) *ExperimentalClient {
11+
return &ExperimentalClient{
12+
Client: client,
13+
}
14+
}

0 commit comments

Comments
 (0)