@@ -115,13 +115,41 @@ type UploadTarFileArgs struct {
115115
116116// WithRecover wraps a HandlerFunc to recover from panics and return an error.
117117func WithRecover [Arg ,Ret any ](h HandlerFunc [Arg ,Ret ])HandlerFunc [Arg ,Ret ] {
118- return func (tb Deps ,args Arg ) (ret Ret ,err error ) {
118+ return func (ctx context. Context , tb Deps ,args Arg ) (ret Ret ,err error ) {
119119defer func () {
120120if r := recover ();r != nil {
121121err = xerrors .Errorf ("tool handler panic: %v" ,r )
122122}
123123}()
124- return h (tb ,args )
124+ return h (ctx ,tb ,args )
125+ }
126+ }
127+
128+ // WithCleanContext wraps a HandlerFunc to provide it with a new context.
129+ // This ensures that no data is passed using context.Value.
130+ // If a deadline is set on the parent context, it will be passed to the child
131+ // context.
132+ func WithCleanContext [Arg ,Ret any ](h HandlerFunc [Arg ,Ret ])HandlerFunc [Arg ,Ret ] {
133+ return func (parent context.Context ,tb Deps ,args Arg ) (ret Ret ,err error ) {
134+ child ,childCancel := context .WithCancel (context .Background ())
135+ defer childCancel ()
136+ // Ensure that cancellation propagates from the parent context to the child context.
137+ go func () {
138+ select {
139+ case <- child .Done ():
140+ return
141+ case <- parent .Done ():
142+ childCancel ()
143+ }
144+ }()
145+ // Also ensure that the child context has the same deadline as the parent
146+ // context.
147+ if deadline ,ok := parent .Deadline ();ok {
148+ deadlineCtx ,deadlineCancel := context .WithDeadline (child ,deadline )
149+ defer deadlineCancel ()
150+ child = deadlineCtx
151+ }
152+ return h (child ,tb ,args )
125153}
126154}
127155
@@ -183,14 +211,14 @@ var (
183211Required : []string {"summary" ,"link" ,"state" },
184212},
185213},
186- Handler :func (tb Deps ,args ReportTaskArgs ) (string ,error ) {
214+ Handler :func (ctx context. Context , tb Deps ,args ReportTaskArgs ) (string ,error ) {
187215if tb .AgentClient == nil {
188216return "" ,xerrors .New ("tool unavailable as CODER_AGENT_TOKEN or CODER_AGENT_TOKEN_FILE not set" )
189217}
190218if tb .AppStatusSlug == "" {
191219return "" ,xerrors .New ("workspace app status slug not found in toolbox" )
192220}
193- if err := tb .AgentClient .PatchAppStatus (context . TODO () , agentsdk.PatchAppStatus {
221+ if err := tb .AgentClient .PatchAppStatus (ctx , agentsdk.PatchAppStatus {
194222AppSlug :tb .AppStatusSlug ,
195223Message :args .Summary ,
196224URI :args .Link ,
@@ -217,12 +245,12 @@ This returns more data than list_workspaces to reduce token usage.`,
217245Required : []string {"workspace_id" },
218246},
219247},
220- Handler :func (tb Deps ,args GetWorkspaceArgs ) (codersdk.Workspace ,error ) {
248+ Handler :func (ctx context. Context , tb Deps ,args GetWorkspaceArgs ) (codersdk.Workspace ,error ) {
221249wsID ,err := uuid .Parse (args .WorkspaceID )
222250if err != nil {
223251return codersdk.Workspace {},xerrors .New ("workspace_id must be a valid UUID" )
224252}
225- return tb .CoderClient .Workspace (context . TODO () ,wsID )
253+ return tb .CoderClient .Workspace (ctx ,wsID )
226254},
227255}
228256
@@ -257,7 +285,7 @@ is provisioned correctly and the agent can connect to the control plane.
257285Required : []string {"user" ,"template_version_id" ,"name" ,"rich_parameters" },
258286},
259287},
260- Handler :func (tb Deps ,args CreateWorkspaceArgs ) (codersdk.Workspace ,error ) {
288+ Handler :func (ctx context. Context , tb Deps ,args CreateWorkspaceArgs ) (codersdk.Workspace ,error ) {
261289tvID ,err := uuid .Parse (args .TemplateVersionID )
262290if err != nil {
263291return codersdk.Workspace {},xerrors .New ("template_version_id must be a valid UUID" )
@@ -272,7 +300,7 @@ is provisioned correctly and the agent can connect to the control plane.
272300Value :v ,
273301})
274302}
275- workspace ,err := tb .CoderClient .CreateUserWorkspace (context . TODO () ,args .User , codersdk.CreateWorkspaceRequest {
303+ workspace ,err := tb .CoderClient .CreateUserWorkspace (ctx ,args .User , codersdk.CreateWorkspaceRequest {
276304TemplateVersionID :tvID ,
277305Name :args .Name ,
278306RichParameterValues :buildParams ,
@@ -297,12 +325,12 @@ is provisioned correctly and the agent can connect to the control plane.
297325},
298326},
299327},
300- Handler :func (tb Deps ,args ListWorkspacesArgs ) ([]MinimalWorkspace ,error ) {
328+ Handler :func (ctx context. Context , tb Deps ,args ListWorkspacesArgs ) ([]MinimalWorkspace ,error ) {
301329owner := args .Owner
302330if owner == "" {
303331owner = codersdk .Me
304332}
305- workspaces ,err := tb .CoderClient .Workspaces (context . TODO () , codersdk.WorkspaceFilter {
333+ workspaces ,err := tb .CoderClient .Workspaces (ctx , codersdk.WorkspaceFilter {
306334Owner :owner ,
307335})
308336if err != nil {
@@ -334,8 +362,8 @@ is provisioned correctly and the agent can connect to the control plane.
334362Required : []string {},
335363},
336364},
337- Handler :func (tb Deps ,_ NoArgs ) ([]MinimalTemplate ,error ) {
338- templates ,err := tb .CoderClient .Templates (context . TODO () , codersdk.TemplateFilter {})
365+ Handler :func (ctx context. Context , tb Deps ,_ NoArgs ) ([]MinimalTemplate ,error ) {
366+ templates ,err := tb .CoderClient .Templates (ctx , codersdk.TemplateFilter {})
339367if err != nil {
340368return nil ,err
341369}
@@ -367,12 +395,12 @@ is provisioned correctly and the agent can connect to the control plane.
367395Required : []string {"template_version_id" },
368396},
369397},
370- Handler :func (tb Deps ,args ListTemplateVersionParametersArgs ) ([]codersdk.TemplateVersionParameter ,error ) {
398+ Handler :func (ctx context. Context , tb Deps ,args ListTemplateVersionParametersArgs ) ([]codersdk.TemplateVersionParameter ,error ) {
371399templateVersionID ,err := uuid .Parse (args .TemplateVersionID )
372400if err != nil {
373401return nil ,xerrors .Errorf ("template_version_id must be a valid UUID: %w" ,err )
374402}
375- parameters ,err := tb .CoderClient .TemplateVersionRichParameters (context . TODO () ,templateVersionID )
403+ parameters ,err := tb .CoderClient .TemplateVersionRichParameters (ctx ,templateVersionID )
376404if err != nil {
377405return nil ,err
378406}
@@ -389,8 +417,8 @@ is provisioned correctly and the agent can connect to the control plane.
389417Required : []string {},
390418},
391419},
392- Handler :func (tb Deps ,_ NoArgs ) (codersdk.User ,error ) {
393- return tb .CoderClient .User (context . TODO () ,"me" )
420+ Handler :func (ctx context. Context , tb Deps ,_ NoArgs ) (codersdk.User ,error ) {
421+ return tb .CoderClient .User (ctx ,"me" )
394422},
395423}
396424
@@ -416,7 +444,7 @@ is provisioned correctly and the agent can connect to the control plane.
416444Required : []string {"workspace_id" ,"transition" },
417445},
418446},
419- Handler :func (tb Deps ,args CreateWorkspaceBuildArgs ) (codersdk.WorkspaceBuild ,error ) {
447+ Handler :func (ctx context. Context , tb Deps ,args CreateWorkspaceBuildArgs ) (codersdk.WorkspaceBuild ,error ) {
420448workspaceID ,err := uuid .Parse (args .WorkspaceID )
421449if err != nil {
422450return codersdk.WorkspaceBuild {},xerrors .Errorf ("workspace_id must be a valid UUID: %w" ,err )
@@ -435,7 +463,7 @@ is provisioned correctly and the agent can connect to the control plane.
435463if templateVersionID != uuid .Nil {
436464cbr .TemplateVersionID = templateVersionID
437465}
438- return tb .CoderClient .CreateWorkspaceBuild (context . TODO () ,workspaceID ,cbr )
466+ return tb .CoderClient .CreateWorkspaceBuild (ctx ,workspaceID ,cbr )
439467},
440468}
441469
@@ -897,8 +925,8 @@ The file_id provided is a reference to a tar file you have uploaded containing t
897925Required : []string {"file_id" },
898926},
899927},
900- Handler :func (tb Deps ,args CreateTemplateVersionArgs ) (codersdk.TemplateVersion ,error ) {
901- me ,err := tb .CoderClient .User (context . TODO () ,"me" )
928+ Handler :func (ctx context. Context , tb Deps ,args CreateTemplateVersionArgs ) (codersdk.TemplateVersion ,error ) {
929+ me ,err := tb .CoderClient .User (ctx ,"me" )
902930if err != nil {
903931return codersdk.TemplateVersion {},err
904932}
@@ -910,7 +938,7 @@ The file_id provided is a reference to a tar file you have uploaded containing t
910938if err != nil {
911939return codersdk.TemplateVersion {},xerrors .Errorf ("template_id must be a valid UUID: %w" ,err )
912940}
913- templateVersion ,err := tb .CoderClient .CreateTemplateVersion (context . TODO () ,me .OrganizationIDs [0 ], codersdk.CreateTemplateVersionRequest {
941+ templateVersion ,err := tb .CoderClient .CreateTemplateVersion (ctx ,me .OrganizationIDs [0 ], codersdk.CreateTemplateVersionRequest {
914942Message :"Created by AI" ,
915943StorageMethod :codersdk .ProvisionerStorageMethodFile ,
916944FileID :fileID ,
@@ -939,12 +967,12 @@ The file_id provided is a reference to a tar file you have uploaded containing t
939967Required : []string {"workspace_agent_id" },
940968},
941969},
942- Handler :func (tb Deps ,args GetWorkspaceAgentLogsArgs ) ([]string ,error ) {
970+ Handler :func (ctx context. Context , tb Deps ,args GetWorkspaceAgentLogsArgs ) ([]string ,error ) {
943971workspaceAgentID ,err := uuid .Parse (args .WorkspaceAgentID )
944972if err != nil {
945973return nil ,xerrors .Errorf ("workspace_agent_id must be a valid UUID: %w" ,err )
946974}
947- logs ,closer ,err := tb .CoderClient .WorkspaceAgentLogsAfter (context . TODO () ,workspaceAgentID ,0 ,false )
975+ logs ,closer ,err := tb .CoderClient .WorkspaceAgentLogsAfter (ctx ,workspaceAgentID ,0 ,false )
948976if err != nil {
949977return nil ,err
950978}
@@ -974,12 +1002,12 @@ The file_id provided is a reference to a tar file you have uploaded containing t
9741002Required : []string {"workspace_build_id" },
9751003},
9761004},
977- Handler :func (tb Deps ,args GetWorkspaceBuildLogsArgs ) ([]string ,error ) {
1005+ Handler :func (ctx context. Context , tb Deps ,args GetWorkspaceBuildLogsArgs ) ([]string ,error ) {
9781006workspaceBuildID ,err := uuid .Parse (args .WorkspaceBuildID )
9791007if err != nil {
9801008return nil ,xerrors .Errorf ("workspace_build_id must be a valid UUID: %w" ,err )
9811009}
982- logs ,closer ,err := tb .CoderClient .WorkspaceBuildLogsAfter (context . TODO () ,workspaceBuildID ,0 )
1010+ logs ,closer ,err := tb .CoderClient .WorkspaceBuildLogsAfter (ctx ,workspaceBuildID ,0 )
9831011if err != nil {
9841012return nil ,err
9851013}
@@ -1005,13 +1033,13 @@ The file_id provided is a reference to a tar file you have uploaded containing t
10051033Required : []string {"template_version_id" },
10061034},
10071035},
1008- Handler :func (tb Deps ,args GetTemplateVersionLogsArgs ) ([]string ,error ) {
1036+ Handler :func (ctx context. Context , tb Deps ,args GetTemplateVersionLogsArgs ) ([]string ,error ) {
10091037templateVersionID ,err := uuid .Parse (args .TemplateVersionID )
10101038if err != nil {
10111039return nil ,xerrors .Errorf ("template_version_id must be a valid UUID: %w" ,err )
10121040}
10131041
1014- logs ,closer ,err := tb .CoderClient .TemplateVersionLogsAfter (context . TODO () ,templateVersionID ,0 )
1042+ logs ,closer ,err := tb .CoderClient .TemplateVersionLogsAfter (ctx ,templateVersionID ,0 )
10151043if err != nil {
10161044return nil ,err
10171045}
@@ -1040,7 +1068,7 @@ The file_id provided is a reference to a tar file you have uploaded containing t
10401068Required : []string {"template_id" ,"template_version_id" },
10411069},
10421070},
1043- Handler :func (tb Deps ,args UpdateTemplateActiveVersionArgs ) (string ,error ) {
1071+ Handler :func (ctx context. Context , tb Deps ,args UpdateTemplateActiveVersionArgs ) (string ,error ) {
10441072templateID ,err := uuid .Parse (args .TemplateID )
10451073if err != nil {
10461074return "" ,xerrors .Errorf ("template_id must be a valid UUID: %w" ,err )
@@ -1049,7 +1077,7 @@ The file_id provided is a reference to a tar file you have uploaded containing t
10491077if err != nil {
10501078return "" ,xerrors .Errorf ("template_version_id must be a valid UUID: %w" ,err )
10511079}
1052- err = tb .CoderClient .UpdateActiveTemplateVersion (context . TODO () ,templateID , codersdk.UpdateActiveTemplateVersion {
1080+ err = tb .CoderClient .UpdateActiveTemplateVersion (ctx ,templateID , codersdk.UpdateActiveTemplateVersion {
10531081ID :templateVersionID ,
10541082})
10551083if err != nil {
@@ -1073,7 +1101,7 @@ The file_id provided is a reference to a tar file you have uploaded containing t
10731101Required : []string {"mime_type" ,"files" },
10741102},
10751103},
1076- Handler :func (tb Deps ,args UploadTarFileArgs ) (codersdk.UploadResponse ,error ) {
1104+ Handler :func (ctx context. Context , tb Deps ,args UploadTarFileArgs ) (codersdk.UploadResponse ,error ) {
10771105pipeReader ,pipeWriter := io .Pipe ()
10781106go func () {
10791107defer pipeWriter .Close ()
@@ -1098,7 +1126,7 @@ The file_id provided is a reference to a tar file you have uploaded containing t
10981126}
10991127}()
11001128
1101- resp ,err := tb .CoderClient .Upload (context . TODO () ,codersdk .ContentTypeTar ,pipeReader )
1129+ resp ,err := tb .CoderClient .Upload (ctx ,codersdk .ContentTypeTar ,pipeReader )
11021130if err != nil {
11031131return codersdk.UploadResponse {},err
11041132}
@@ -1133,16 +1161,16 @@ The file_id provided is a reference to a tar file you have uploaded containing t
11331161Required : []string {"name" ,"display_name" ,"description" ,"version_id" },
11341162},
11351163},
1136- Handler :func (tb Deps ,args CreateTemplateArgs ) (codersdk.Template ,error ) {
1137- me ,err := tb .CoderClient .User (context . TODO () ,"me" )
1164+ Handler :func (ctx context. Context , tb Deps ,args CreateTemplateArgs ) (codersdk.Template ,error ) {
1165+ me ,err := tb .CoderClient .User (ctx ,"me" )
11381166if err != nil {
11391167return codersdk.Template {},err
11401168}
11411169versionID ,err := uuid .Parse (args .VersionID )
11421170if err != nil {
11431171return codersdk.Template {},xerrors .Errorf ("version_id must be a valid UUID: %w" ,err )
11441172}
1145- template ,err := tb .CoderClient .CreateTemplate (context . TODO () ,me .OrganizationIDs [0 ], codersdk.CreateTemplateRequest {
1173+ template ,err := tb .CoderClient .CreateTemplate (ctx ,me .OrganizationIDs [0 ], codersdk.CreateTemplateRequest {
11461174Name :args .Name ,
11471175DisplayName :args .DisplayName ,
11481176Description :args .Description ,
@@ -1167,12 +1195,12 @@ The file_id provided is a reference to a tar file you have uploaded containing t
11671195},
11681196},
11691197},
1170- Handler :func (tb Deps ,args DeleteTemplateArgs ) (string ,error ) {
1198+ Handler :func (ctx context. Context , tb Deps ,args DeleteTemplateArgs ) (string ,error ) {
11711199templateID ,err := uuid .Parse (args .TemplateID )
11721200if err != nil {
11731201return "" ,xerrors .Errorf ("template_id must be a valid UUID: %w" ,err )
11741202}
1175- err = tb .CoderClient .DeleteTemplate (context . TODO () ,templateID )
1203+ err = tb .CoderClient .DeleteTemplate (ctx ,templateID )
11761204if err != nil {
11771205return "" ,err
11781206}