diff --git a/cli/azd/extensions/azure.ai.agents/internal/project/parser.go b/cli/azd/extensions/azure.ai.agents/internal/project/parser.go index fd0687f3e40..c73d8b0e60b 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/project/parser.go +++ b/cli/azd/extensions/azure.ai.agents/internal/project/parser.go @@ -146,7 +146,7 @@ func (p *FoundryParser) SetIdentity(ctx context.Context, args *azdext.ProjectEve // Get Application ID from Principal ID fmt.Println("Retrieving Application ID...") - projectClientID, err := getApplicationID(context.Background(), cred, projectPrincipalID) + projectClientID, err := getApplicationID(ctx, cred, projectPrincipalID) if err != nil { return fmt.Errorf("failed to get Application ID: %w", err) } @@ -484,12 +484,12 @@ func (p *FoundryParser) CoboPostDeploy(ctx context.Context, args *azdext.Project } // Register agent with retry logic - agentVersion := registerAgent(uri, token, resourceID, ingressSuffix) + agentVersion := registerAgent(ctx, uri, token, resourceID, ingressSuffix) // Test authentication and agent if agentVersion != "" { - testUnauthenticatedAccess(acaEndpoint) - testDataPlane(aiFoundryProjectEndpoint, token, agentName, agentVersion) + testUnauthenticatedAccess(ctx, acaEndpoint) + testDataPlane(ctx, aiFoundryProjectEndpoint, token, agentName, agentVersion) } // Print Azure Portal link @@ -1005,7 +1005,7 @@ func getLatestRevisionName( } // registerAgent registers the agent with Microsoft Foundry -func registerAgent(uri, token, resourceID, ingressSuffix string) string { +func registerAgent(ctx context.Context, uri, token, resourceID, ingressSuffix string) string { fmt.Println() fmt.Println("======================================") fmt.Println("Registering Agent Version") @@ -1038,7 +1038,7 @@ func registerAgent(uri, token, resourceID, ingressSuffix string) string { } client := &http.Client{} - req, err := http.NewRequest("POST", uri, bytes.NewBuffer(payloadBytes)) + req, err := http.NewRequestWithContext(ctx, "POST", uri, bytes.NewBuffer(payloadBytes)) if err != nil { fmt.Fprintf(os.Stderr, "Error creating request: %v\n", err) continue @@ -1090,7 +1090,7 @@ func registerAgent(uri, token, resourceID, ingressSuffix string) string { } // testUnauthenticatedAccess tests unauthenticated access (should return 401) -func testUnauthenticatedAccess(acaEndpoint string) { +func testUnauthenticatedAccess(ctx context.Context, acaEndpoint string) { fmt.Println() fmt.Println("======================================") fmt.Println("Testing Unauthenticated Access") @@ -1103,7 +1103,7 @@ func testUnauthenticatedAccess(acaEndpoint string) { fmt.Printf("Request Body: %s\n", string(payload)) client := &http.Client{} - req, err := http.NewRequest("POST", uri, bytes.NewBuffer(payload)) + req, err := http.NewRequestWithContext(ctx, "POST", uri, bytes.NewBuffer(payload)) if err != nil { fmt.Fprintf(os.Stderr, "Error creating request: %v\n", err) return @@ -1134,7 +1134,7 @@ func testUnauthenticatedAccess(acaEndpoint string) { } // testDataPlane tests the agent data plane with authenticated request -func testDataPlane(endpoint, token, agentName, agentVersion string) { +func testDataPlane(ctx context.Context, endpoint, token, agentName, agentVersion string) { fmt.Println() fmt.Println("======================================") fmt.Println("Testing Agent Data Plane") @@ -1158,7 +1158,7 @@ func testDataPlane(endpoint, token, agentName, agentVersion string) { fmt.Println(string(payloadBytes)) client := &http.Client{} - req, err := http.NewRequest("POST", uri, bytes.NewBuffer(payloadBytes)) + req, err := http.NewRequestWithContext(ctx, "POST", uri, bytes.NewBuffer(payloadBytes)) if err != nil { fmt.Fprintf(os.Stderr, "Error creating request: %v\n", err) return diff --git a/cli/azd/internal/telemetry/appinsights-exporter/transmitter.go b/cli/azd/internal/telemetry/appinsights-exporter/transmitter.go index 772b07d0714..1426c954fdf 100644 --- a/cli/azd/internal/telemetry/appinsights-exporter/transmitter.go +++ b/cli/azd/internal/telemetry/appinsights-exporter/transmitter.go @@ -8,6 +8,7 @@ package appinsightsexporter import ( "bytes" "compress/gzip" + "context" "encoding/json" "io" "net/http" @@ -16,7 +17,7 @@ import ( ) type Transmitter interface { - Transmit(payload []byte, items TelemetryItems) (*TransmissionResult, error) + Transmit(ctx context.Context, payload []byte, items TelemetryItems) (*TransmissionResult, error) } type httpTransmitter struct { @@ -63,7 +64,11 @@ func NewTransmitter(endpointAddress string, client *http.Client) Transmitter { return &httpTransmitter{endpointAddress, client} } -func (transmitter *httpTransmitter) Transmit(payload []byte, items TelemetryItems) (*TransmissionResult, error) { +func (transmitter *httpTransmitter) Transmit( + ctx context.Context, + payload []byte, + items TelemetryItems, +) (*TransmissionResult, error) { startTime := time.Now() // Compress the payload @@ -77,7 +82,7 @@ func (transmitter *httpTransmitter) Transmit(payload []byte, items TelemetryItem gzipWriter.Close() - req, err := http.NewRequest(http.MethodPost, transmitter.endpoint, &postBody) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, transmitter.endpoint, &postBody) if err != nil { return nil, err } diff --git a/cli/azd/internal/telemetry/appinsights-exporter/transmitter_test.go b/cli/azd/internal/telemetry/appinsights-exporter/transmitter_test.go index 00519f57124..265a2c6113a 100644 --- a/cli/azd/internal/telemetry/appinsights-exporter/transmitter_test.go +++ b/cli/azd/internal/telemetry/appinsights-exporter/transmitter_test.go @@ -6,6 +6,7 @@ package appinsightsexporter import ( "bytes" "compress/gzip" + "context" "fmt" "io" "net/http" @@ -106,7 +107,7 @@ func doBasicTransmit(client Transmitter, server *testServer, t *testing.T) { server.responseData = []byte(`{"itemsReceived":3, "itemsAccepted":5, "errors":[]}`) server.responseHeaders["Content-type"] = "application/json" - result, err := client.Transmit([]byte("foobar"), make(TelemetryItems, 0)) + result, err := client.Transmit(context.Background(), []byte("foobar"), make(TelemetryItems, 0)) if err != nil { t.Log(err.Error()) } @@ -185,7 +186,7 @@ func TestFailedTransmit(t *testing.T) { `{"itemsReceived":3, "itemsAccepted":0, "errors":[{"index": 2, "statusCode": 500, "message": "Hello"}]}`, ) server.responseHeaders["Content-type"] = "application/json" - result, err := client.Transmit([]byte("foobar"), make(TelemetryItems, 0)) + result, err := client.Transmit(context.Background(), []byte("foobar"), make(TelemetryItems, 0)) server.waitForRequest(t) if err != nil { @@ -237,7 +238,7 @@ func TestThrottledTransmit(t *testing.T) { server.responseData = make([]byte, 0) server.responseHeaders["Content-type"] = "application/json" server.responseHeaders["retry-after"] = "Wed, 09 Aug 2017 23:43:57 UTC" - result, err := client.Transmit([]byte("foobar"), make(TelemetryItems, 0)) + result, err := client.Transmit(context.Background(), []byte("foobar"), make(TelemetryItems, 0)) server.waitForRequest(t) if err != nil { diff --git a/cli/azd/internal/telemetry/uploader.go b/cli/azd/internal/telemetry/uploader.go index 519e5d9d1c0..b3b276fcbc1 100644 --- a/cli/azd/internal/telemetry/uploader.go +++ b/cli/azd/internal/telemetry/uploader.go @@ -180,7 +180,7 @@ func (u *TelemetryUploader) transmit(ctx context.Context, item *StoredItem) { // When in debug mode, we deserialize to get better error messages telemetryItems.Deserialize(payload) } - result, err := u.transmitter.Transmit(payload, telemetryItems) + result, err := u.transmitter.Transmit(ctx, payload, telemetryItems) if err == nil && result != nil && result.IsSuccess() { return } diff --git a/cli/azd/internal/telemetry/uploader_test.go b/cli/azd/internal/telemetry/uploader_test.go index 080868f490f..388947d9772 100644 --- a/cli/azd/internal/telemetry/uploader_test.go +++ b/cli/azd/internal/telemetry/uploader_test.go @@ -116,6 +116,7 @@ func NewTransmitterStub() *TransmitterStub { } func (tr *TransmitterStub) Transmit( + _ context.Context, payload []byte, items appinsightsexporter.TelemetryItems, ) (*appinsightsexporter.TransmissionResult, error) { diff --git a/cli/azd/pkg/llm/github_copilot.go b/cli/azd/pkg/llm/github_copilot.go index 809b9db4093..de0d343b788 100644 --- a/cli/azd/pkg/llm/github_copilot.go +++ b/cli/azd/pkg/llm/github_copilot.go @@ -258,7 +258,7 @@ func copilotToken(ctx context.Context, console input.Console) (*tokenData, error } // Token is missing or expired, get a new one - newToken, err := newCopilotToken(githubToken) + newToken, err := newCopilotToken(ctx, githubToken) if err != nil { // If Copilot token request fails, GitHub token might be expired if strings.Contains(err.Error(), "status 401") || strings.Contains(err.Error(), "status 403") { @@ -281,7 +281,7 @@ func copilotToken(ctx context.Context, console input.Console) (*tokenData, error } // Try getting Copilot token again with new GitHub token - newToken, err = newCopilotToken(githubToken) + newToken, err = newCopilotToken(ctx, githubToken) if err != nil { return nil, fmt.Errorf("failed to get Copilot token after re-authentication: %w", err) } @@ -383,9 +383,14 @@ func isTokenExpired(expiresAt int64) bool { } // newCopilotToken gets a Copilot session token using the GitHub token -func newCopilotToken(githubToken string) (*tokenData, error) { +func newCopilotToken(ctx context.Context, githubToken string) (*tokenData, error) { client := &http.Client{} - req, err := http.NewRequest("GET", "https://api.github.com/copilot_internal/v2/token", nil) + req, err := http.NewRequestWithContext( + ctx, + "GET", + "https://api.github.com/copilot_internal/v2/token", + nil, + ) if err != nil { return nil, err } diff --git a/cli/azd/tools/avmres/main.go b/cli/azd/tools/avmres/main.go index 0c363776f4b..b7d1a465d5d 100644 --- a/cli/azd/tools/avmres/main.go +++ b/cli/azd/tools/avmres/main.go @@ -4,6 +4,7 @@ package main import ( + "context" "fmt" "io" "net/http" @@ -190,7 +191,7 @@ func fetchGithub( repo string, path string, tag string) ([]byte, error) { - req, err := http.NewRequest("GET", + req, err := http.NewRequestWithContext(context.Background(), "GET", fmt.Sprintf("https://api.github.com/repos/%s/contents/%s?ref=%s", repo, path, tag), nil) if err != nil { return nil, err