Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 10 additions & 10 deletions cli/azd/extensions/azure.ai.agents/internal/project/parser.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand All @@ -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
Expand Down Expand Up @@ -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")
Expand All @@ -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
Expand Down
11 changes: 8 additions & 3 deletions cli/azd/internal/telemetry/appinsights-exporter/transmitter.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ package appinsightsexporter
import (
"bytes"
"compress/gzip"
"context"
"encoding/json"
"io"
"net/http"
Expand All @@ -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 {
Expand Down Expand Up @@ -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
Expand All @@ -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
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ package appinsightsexporter
import (
"bytes"
"compress/gzip"
"context"
"fmt"
"io"
"net/http"
Expand Down Expand Up @@ -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())
}
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down
2 changes: 1 addition & 1 deletion cli/azd/internal/telemetry/uploader.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down
1 change: 1 addition & 0 deletions cli/azd/internal/telemetry/uploader_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,7 @@ func NewTransmitterStub() *TransmitterStub {
}

func (tr *TransmitterStub) Transmit(
_ context.Context,
payload []byte,
items appinsightsexporter.TelemetryItems,
) (*appinsightsexporter.TransmissionResult, error) {
Expand Down
13 changes: 9 additions & 4 deletions cli/azd/pkg/llm/github_copilot.go
Original file line number Diff line number Diff line change
Expand Up @@ -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") {
Expand All @@ -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)
}
Expand Down Expand Up @@ -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
}
Expand Down
3 changes: 2 additions & 1 deletion cli/azd/tools/avmres/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
package main

import (
"context"
"fmt"
"io"
"net/http"
Expand Down Expand Up @@ -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
Expand Down
Loading