From 69a8b0bbfe9d81e32ac243673b06100c8af8a044 Mon Sep 17 00:00:00 2001 From: Wei Meng Date: Wed, 9 Sep 2026 14:39:29 +0800 Subject: [PATCH 1/5] feat(ai): support Invocations protocol in shared lifecycle commands --- .../docs/specs/long-running-agent-invoke.md | 2 +- .../internal/cmd/agent_endpoint.go | 24 ++ .../azure.ai.agents/internal/cmd/delete.go | 2 +- .../internal/cmd/invocations.go | 20 +- .../internal/cmd/invocations_test.go | 27 +- .../azure.ai.agents/internal/cmd/invoke.go | 22 +- .../internal/cmd/invoke_invocation.go | 348 ++++++++++++++++++ .../internal/cmd/invoke_invocation_test.go | 150 ++++++++ .../internal/cmd/invoke_response_test.go | 9 + .../azure.ai.agents/internal/cmd/listen.go | 10 +- .../internal/exterrors/codes.go | 20 +- 11 files changed, 603 insertions(+), 31 deletions(-) create mode 100644 cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go create mode 100644 cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go diff --git a/cli/azd/extensions/azure.ai.agents/docs/specs/long-running-agent-invoke.md b/cli/azd/extensions/azure.ai.agents/docs/specs/long-running-agent-invoke.md index 59893108b0e..6ec97dbb59b 100644 --- a/cli/azd/extensions/azure.ai.agents/docs/specs/long-running-agent-invoke.md +++ b/cli/azd/extensions/azure.ai.agents/docs/specs/long-running-agent-invoke.md @@ -42,7 +42,7 @@ azd ai agent invocations show --protocol invocations --id azd ai agent invocations cancel --protocol invocations --id ``` -The Invocations-protocol lifecycle implementation is delivered in the second PR. Its existing synchronous, SSE, raw, and `202 Accepted` polling behavior on create is unchanged. +The Invocations-protocol lifecycle implementation is included in PR #9901. Its existing synchronous, SSE, raw, and `202 Accepted` polling behavior on create is unchanged. Creation captures the ID from `x-agent-invocation-id`, or from the `invocation_id` property of a `202 Accepted` body. The body is preserved for the existing polling and raw-output handlers. Show and cancel use the resolved endpoint/API version, without inheriting session context from the create. ### Protocol selection diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/agent_endpoint.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/agent_endpoint.go index c1f6ab9a275..926a48603a0 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/agent_endpoint.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/agent_endpoint.go @@ -198,6 +198,30 @@ func buildInvocationsURL(projectEndpoint, agentName, apiVersion, sid string) str return invURL } +func buildInvocationLifecycleURL( + projectEndpoint string, + agentName string, + invocationID string, + apiVersion string, +) string { + if apiVersion == "" { + apiVersion = DefaultAgentAPIVersion + } + return fmt.Sprintf( + "%s/agents/%s/endpoint/protocols/invocations/%s?api-version=%s", + projectEndpoint, + agentName, + url.PathEscape(invocationID), + url.QueryEscape(apiVersion), + ) +} + +func buildInvocationCancelURL(projectEndpoint, agentName, invocationID, apiVersion string) string { + lifecycleURL := buildInvocationLifecycleURL(projectEndpoint, agentName, invocationID, apiVersion) + parts := strings.SplitN(lifecycleURL, "?", 2) + return parts[0] + "/cancel?" + parts[1] +} + // buildA2AInvokeURL builds the Foundry "a2a" protocol URL for an agent. When sid // is non-empty, an agent_session_id query parameter is appended (URL-encoded) so // the request routes to the same agent session, matching the invocations protocol. diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/delete.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/delete.go index 9c884804e55..579b3380cbd 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/delete.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/delete.go @@ -175,7 +175,7 @@ func (a *DeleteAction) Run(ctx context.Context) error { return classifyDeleteError(err, agentName) } - // Best-effort: clean up saved session, conversation, and current Response state (same as postdown hook). + // Best-effort: clean up saved session, conversation, Response, and Invocation state (same as postdown hook). // Must run before cleanupEnvVars since it reads AGENT_{KEY}_ENDPOINT. if envResp, err := azdClient.Environment().GetCurrent(ctx, &azdext.EmptyRequest{}); err == nil { cleanupAgentState(ctx, azdClient, envResp.Environment.Name, info.ServiceName) diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invocations.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invocations.go index 16b8e83be71..50db424f14c 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/invocations.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/invocations.go @@ -141,8 +141,14 @@ func validateInvocationCommandFlags(cmd *cobra.Command, flags *invocationCommand // supportsInvocationOperation describes implemented CLI operations, not the // capabilities of every deployed agent. The service may still reject a request. func supportsInvocationOperation(protocol agent_api.AgentProtocol, operation invocationOperation) bool { - return protocol == agent_api.AgentProtocolResponses && - (operation == invocationShow || operation == invocationFollow || operation == invocationCancel) + switch protocol { + case agent_api.AgentProtocolResponses: + return operation == invocationShow || operation == invocationFollow || operation == invocationCancel + case agent_api.AgentProtocolInvocations: + return operation == invocationShow || operation == invocationCancel + default: + return false + } } // resolveInvocationCommand resolves a protocol before reading its current ID. @@ -215,6 +221,14 @@ func resolveCurrentInvocationID( if record != nil { id = record.ResponseID } + case agent_api.AgentProtocolInvocations: + record, err := newInvocationStateStore(rc.azdClient).Get(ctx, rc.agentKey) + if err != nil { + return "", classifyInvocationStateReadError(err) + } + if record != nil { + id = record.InvocationID + } } if id == "" { return "", exterrors.Validation(exterrors.CodeInvalidParameter, @@ -253,6 +267,8 @@ func (a *InvokeAction) runInvocationOperation( switch agent_api.AgentProtocol(a.flags.protocol) { case agent_api.AgentProtocolResponses: return a.runResponseOperation(ctx, rc, id, operation, format, writer) + case agent_api.AgentProtocolInvocations: + return a.runInvocationsProtocolOperation(ctx, rc, id, operation, format, writer) default: return exterrors.Validation(exterrors.CodeInvalidParameter, "unsupported invocation protocol", "select a protocol that supports this operation") diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invocations_test.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invocations_test.go index 0d03d6cd19e..57832cdc1b8 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/invocations_test.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/invocations_test.go @@ -203,21 +203,29 @@ func TestResolveInvocationCommandSelection(t *testing.T) { go func() { _ = server.Serve(listener) }() t.Cleanup(server.Stop) t.Setenv("AZD_SERVER", listener.Addr().String()) - const endpoint = "https://example.services.ai.azure.com/api/projects/project/agents/agent/" + - "endpoint/protocols/openai/responses?api-version=v1" + const endpoint = "https://example.services.ai.azure.com/api/projects/project/agents/agent/endpoint/protocols/" key := buildAgentKey("https://example.services.ai.azure.com/api/projects/project", "agent", "", false) config.setJSON(t, responsesConfigPath, map[string]savedResponse{key: {ResponseID: "resp_current"}}) - for _, tt := range []struct{ name, explicitID, want string }{ - {name: "implicit", want: "resp_current"}, - {name: "explicit", explicitID: "resp_explicit", want: "resp_explicit"}, + config.setJSON(t, invocationsConfigPath, map[string]savedInvocation{key: {InvocationID: "inv_current"}}) + for _, tt := range []struct{ name, path, protocol, explicitID, want string }{ + {name: "Responses implicit", path: "openai/responses", protocol: "responses", want: "resp_current"}, + {name: "Responses explicit", path: "openai/responses", protocol: "responses", + explicitID: "resp_explicit", want: "resp_explicit"}, + {name: "Invocations implicit", path: "invocations", protocol: "invocations", want: "inv_current"}, + {name: "Invocations explicit", path: "invocations", protocol: "invocations", + explicitID: "inv_explicit", want: "inv_explicit"}, } { t.Run(tt.name, func(t *testing.T) { - _, rc, id, err := resolveInvocationCommand(t.Context(), &invocationCommandFlags{ - agentEndpoint: endpoint, id: tt.explicitID, + action, rc, id, err := resolveInvocationCommand(t.Context(), &invocationCommandFlags{ + agentEndpoint: endpoint + tt.path + "?api-version=v1", id: tt.explicitID, }, invocationShow) require.NoError(t, err) defer rc.azdClient.Close() assert.Equal(t, tt.want, id) + assert.Equal(t, tt.protocol, action.flags.protocol) + var invocations map[string]savedInvocation + config.getJSON(t, invocationsConfigPath, &invocations) + assert.Equal(t, "inv_current", invocations[key].InvocationID) var saved map[string]savedResponse config.getJSON(t, responsesConfigPath, &saved) assert.Equal(t, "resp_current", saved[key].ResponseID) @@ -231,8 +239,9 @@ func TestInvocationOperationSupport(t *testing.T) { "activity", "invocations_ws", "voice", } { for _, operation := range []invocationOperation{invocationShow, invocationFollow, invocationCancel} { - assert.Equal(t, protocol == agent_api.AgentProtocolResponses, - supportsInvocationOperation(protocol, operation), "%s %s", protocol, operation) + want := protocol == agent_api.AgentProtocolResponses || + (protocol == agent_api.AgentProtocolInvocations && operation != invocationFollow) + assert.Equal(t, want, supportsInvocationOperation(protocol, operation), "%s %s", protocol, operation) } } } diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke.go index ccdea6bd0fd..653075749cf 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke.go @@ -1865,16 +1865,22 @@ func (a *InvokeAction) invocationsRemote(ctx context.Context) error { ttfb := time.Since(invokeStart) defer resp.Body.Close() - // Print the invocation ID if the agent returned one. We do not persist it - // to the per-user config: the config store only supports the "sessions" - // and "conversations" maps (see validateStoreField), and invocation IDs - // are not used to drive any subsequent invoke -- they are emitted purely - // for trace correlation. - if !raw { - if invID := resp.Header.Get("x-agent-invocation-id"); invID != "" { - fmt.Printf("Invocation: %s\n", invID) + invocationID, err := invocationIDFromResponse(resp) + if err != nil { + return err + } + if resp.StatusCode < http.StatusBadRequest && invocationID != "" && rc.azdClient != nil && agentKey != "" { + if err := newInvocationStateStore(rc.azdClient).Save( + ctx, + agentKey, + savedInvocation{InvocationID: invocationID}, + ); err != nil { + fmt.Fprintf(os.Stderr, "warning: Invocation %s was accepted, but its ID was not saved: %v\n", invocationID, err) } } + if !raw && invocationID != "" { + fmt.Printf("Invocation: %s\n", invocationID) + } // Always capture session state from response headers (needed even in raw mode // so subsequent invokes can reuse the session). Reads headers, not the body. diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go new file mode 100644 index 00000000000..bee7baa0d65 --- /dev/null +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go @@ -0,0 +1,348 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package cmd + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "strconv" + "text/tabwriter" + "time" + + "azureaiagent/internal/exterrors" + + "github.com/azure/azure-dev/cli/azd/pkg/azdext" +) + +func invocationIDFromResponse(resp *http.Response) (string, error) { + if invocationID := resp.Header.Get("x-agent-invocation-id"); invocationID != "" { + return invocationID, nil + } + if resp.StatusCode != http.StatusAccepted { + return "", nil + } + + originalBody := resp.Body + body, err := io.ReadAll(originalBody) + _ = originalBody.Close() + if err != nil { + return "", fmt.Errorf("read accepted Invocation response: %w", err) + } + resp.Body = io.NopCloser(bytes.NewReader(body)) + var accepted struct { + InvocationID string `json:"invocation_id"` + } + if err := json.Unmarshal(body, &accepted); err != nil { + return "", nil + } + return accepted.InvocationID, nil +} + +type invocationSnapshot struct { + ID string `json:"id"` + InvocationID string `json:"invocation_id"` + Status string `json:"status"` +} + +type invocationSnapshotResult struct { + snapshot invocationSnapshot + raw []byte +} + +func (a *InvokeAction) runInvocationsProtocolOperation( + ctx context.Context, + rc *remoteContext, + id string, + operation invocationOperation, + format string, + writer io.Writer, +) error { + switch operation { + case invocationShow: + result, err := a.getInvocation(ctx, rc, id) + if err != nil { + return classifyInvocationLifecycleError(err, exterrors.OpShowInvocation, "showing Invocation") + } + return printInvocationSnapshot(writer, result, format) + case invocationCancel: + return classifyInvocationLifecycleError( + a.cancelInvocation(ctx, rc, id, writer), exterrors.OpCancelInvocation, "cancelling Invocation", + ) + default: + return exterrors.Validation(exterrors.CodeInvalidParameter, + fmt.Sprintf("invocations %s is not supported with the invocations protocol", operation), + "use invocations show to retrieve a snapshot") + } +} + +func classifyInvocationStateReadError(cause error) error { + if _, ok := errors.AsType[*azdext.ConfigError](cause); !ok { + return exterrors.FromHost(cause, exterrors.OpReadInvocationState, "reading current Invocation state failed") + } + return exterrors.Validation( + exterrors.CodeInvalidInvocationState, + fmt.Sprintf("saved Invocation state at %q could not be read: %v", invocationsConfigPath, cause), + fmt.Sprintf("clear the invalid state with `azd config unset %s`, or repair that config value", invocationsConfigPath), + ) +} + +func (a *InvokeAction) cancelInvocation( + ctx context.Context, rc *remoteContext, invocationID string, writer io.Writer, +) error { + token, err := a.acquireBearerToken(ctx) + if err != nil { + return err + } + cancelURL := buildInvocationCancelURL(rc.projectEndpoint, rc.name, invocationID, rc.apiVersion) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, cancelURL, nil) + if err != nil { + return fmt.Errorf("create Invocation cancel request: %w", err) + } + req.Header.Set("Authorization", "Bearer "+token) + applyCustomHeaders(req, a.clientHeaders) + applyRemoteUserIdentityHeader(req, &a.flags.userIdentityFlags) + + //nolint:gosec // URL is built from a validated Foundry endpoint. + resp, err := (&http.Client{Timeout: 30 * time.Second}).Do(req) + if err != nil { + return fmt.Errorf("cancel Invocation %s: %w", invocationID, err) + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + return fmt.Errorf("read Invocation cancel result: %w", err) + } + if resp.StatusCode >= http.StatusBadRequest { + result, getErr := a.getInvocation(ctx, rc, invocationID) + if getErr == nil && isTerminalInvocationStatus(result.snapshot.Status) { + _, err = fmt.Fprintf( + writer, + "Invocation %s is already %s; nothing to cancel.\n", + invocationID, + result.snapshot.Status, + ) + return err + } + return &invocationLifecycleHTTPError{ + method: http.MethodPost, + requestURL: cancelURL, + statusCode: resp.StatusCode, + status: resp.Status, + body: body, + } + } + + var result invocationSnapshot + if json.Unmarshal(body, &result) == nil && result.Status != "" { + _, err = fmt.Fprintf(writer, "Invocation %s is %s.\n", invocationID, result.Status) + return err + } + _, err = fmt.Fprintf(writer, "Cancellation requested for Invocation %s.\n", invocationID) + return err +} + +func (a *InvokeAction) getInvocation( + ctx context.Context, + rc *remoteContext, + invocationID string, +) (invocationSnapshotResult, error) { + token, err := a.acquireBearerToken(ctx) + if err != nil { + return invocationSnapshotResult{}, err + } + requestURL := buildInvocationLifecycleURL(rc.projectEndpoint, rc.name, invocationID, rc.apiVersion) + req, err := http.NewRequestWithContext(ctx, http.MethodGet, requestURL, nil) + if err != nil { + return invocationSnapshotResult{}, fmt.Errorf("create Invocation show request: %w", err) + } + req.Header.Set("Authorization", "Bearer "+token) + applyCustomHeaders(req, a.clientHeaders) + applyRemoteUserIdentityHeader(req, &a.flags.userIdentityFlags) + + //nolint:gosec // URL is built from a validated Foundry endpoint. + resp, err := (&http.Client{Timeout: 30 * time.Second}).Do(req) + if err != nil { + return invocationSnapshotResult{}, fmt.Errorf("show Invocation %s: %w", invocationID, err) + } + defer resp.Body.Close() + body, err := io.ReadAll(resp.Body) + if err != nil { + return invocationSnapshotResult{}, fmt.Errorf("read Invocation: %w", err) + } + if resp.StatusCode >= http.StatusBadRequest { + return invocationSnapshotResult{}, &invocationLifecycleHTTPError{ + method: http.MethodGet, + requestURL: requestURL, + statusCode: resp.StatusCode, + status: resp.Status, + body: body, + } + } + var snapshot invocationSnapshot + if err := json.Unmarshal(body, &snapshot); err != nil { + return invocationSnapshotResult{}, fmt.Errorf("decode Invocation: %w", err) + } + actualID := snapshot.ID + if actualID == "" { + actualID = snapshot.InvocationID + } + if actualID != "" && actualID != invocationID { + return invocationSnapshotResult{}, fmt.Errorf( + "Invocation ID %q does not match requested ID %q", + actualID, + invocationID, + ) + } + return invocationSnapshotResult{snapshot: snapshot, raw: body}, nil +} + +type invocationLifecycleHTTPError struct { + method string + requestURL string + statusCode int + status string + body []byte +} + +func (e *invocationLifecycleHTTPError) Error() string { + return fmt.Sprintf("%s %s failed with HTTP %d: %s\n%s", e.method, e.requestURL, e.statusCode, e.status, e.body) +} + +func classifyInvocationLifecycleError(cause error, operation, label string) error { + if cause == nil { + return nil + } + httpErr, ok := errors.AsType[*invocationLifecycleHTTPError](cause) + if !ok { + return cause + } + serviceName := "" + if parsed, err := url.Parse(httpErr.requestURL); err == nil { + serviceName = parsed.Hostname() + } + serviceErr := exterrors.Service( + operation, + strconv.Itoa(httpErr.statusCode), + fmt.Sprintf("%s failed with HTTP %d: %s", label, httpErr.statusCode, httpErr.status), + serviceName, + "", + ) + serviceErr.StatusCode = httpErr.statusCode + return serviceErr +} + +func printInvocationSnapshot(writer io.Writer, result invocationSnapshotResult, format string) error { + if format != "table" { + var formatted any + if err := json.Unmarshal(result.raw, &formatted); err != nil { + return fmt.Errorf("decode Invocation JSON: %w", err) + } + data, err := json.MarshalIndent(formatted, "", " ") + if err != nil { + return fmt.Errorf("format Invocation JSON: %w", err) + } + _, err = fmt.Fprintln(writer, string(data)) + return err + } + + id := result.snapshot.ID + if id == "" { + id = result.snapshot.InvocationID + } + table := tabwriter.NewWriter(writer, 0, 0, 2, ' ', 0) + fmt.Fprintln(table, "FIELD\tVALUE") + fmt.Fprintln(table, "-----\t-----") + fmt.Fprintf(table, "Invocation ID\t%s\n", id) + fmt.Fprintf(table, "Status\t%s\n", result.snapshot.Status) + return table.Flush() +} + +func isTerminalInvocationStatus(status string) bool { + switch status { + case "completed", "failed", "cancelled", "canceled": + return true + default: + return false + } +} + +const invocationsConfigPath = configPathPrefix + ".invocations" + +type savedInvocation struct { + InvocationID string `json:"invocationId"` +} + +type invocationStateStore struct { + client *azdext.AzdClient +} + +func newInvocationStateStore(client *azdext.AzdClient) *invocationStateStore { + return &invocationStateStore{client: client} +} + +func (s *invocationStateStore) Get(ctx context.Context, agentKey string) (*savedInvocation, error) { + config, err := azdext.NewConfigHelper(s.client) + if err != nil { + return nil, fmt.Errorf("create invocation config helper: %w", err) + } + var records map[string]savedInvocation + found, err := config.GetUserJSON(ctx, invocationsConfigPath, &records) + if err != nil { + return nil, fmt.Errorf("read invocations: %w", err) + } + if !found || records == nil { + return nil, nil + } + record, ok := records[agentKey] + if !ok { + return nil, nil + } + return &record, nil +} + +func (s *invocationStateStore) Save(ctx context.Context, agentKey string, record savedInvocation) error { + config, err := azdext.NewConfigHelper(s.client) + if err != nil { + return fmt.Errorf("create invocation config helper: %w", err) + } + var records map[string]savedInvocation + found, err := config.GetUserJSON(ctx, invocationsConfigPath, &records) + if err != nil { + return fmt.Errorf("read invocations: %w", err) + } + if !found || records == nil { + records = make(map[string]savedInvocation) + } + records[agentKey] = record + if err := config.SetUserJSON(ctx, invocationsConfigPath, records); err != nil { + return fmt.Errorf("write invocations: %w", err) + } + return nil +} + +func (s *invocationStateStore) Delete(ctx context.Context, agentKey string) error { + config, err := azdext.NewConfigHelper(s.client) + if err != nil { + return fmt.Errorf("create invocation config helper: %w", err) + } + var records map[string]savedInvocation + found, err := config.GetUserJSON(ctx, invocationsConfigPath, &records) + if err != nil { + return fmt.Errorf("read invocations: %w", err) + } + if !found || records == nil { + return nil + } + delete(records, agentKey) + if err := config.SetUserJSON(ctx, invocationsConfigPath, records); err != nil { + return fmt.Errorf("write invocations: %w", err) + } + return nil +} diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go new file mode 100644 index 00000000000..696af8f5f64 --- /dev/null +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go @@ -0,0 +1,150 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +package cmd + +import ( + "bytes" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type trackingReadCloser struct { + io.Reader + closed bool +} + +func (r *trackingReadCloser) Close() error { + r.closed = true + return nil +} + +func TestInvocationIDFromResponse(t *testing.T) { + t.Run("header", func(t *testing.T) { + resp := &http.Response{Header: http.Header{"X-Agent-Invocation-Id": []string{"inv_header"}}} + id, err := invocationIDFromResponse(resp) + require.NoError(t, err) + assert.Equal(t, "inv_header", id) + }) + + t.Run("successful body is untouched", func(t *testing.T) { + body := `{"result":"ok"}` + original := &trackingReadCloser{Reader: strings.NewReader(body)} + resp := &http.Response{StatusCode: http.StatusOK, Header: make(http.Header), Body: original} + id, err := invocationIDFromResponse(resp) + require.NoError(t, err) + assert.Empty(t, id) + assert.False(t, original.closed) + remaining, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Equal(t, body, string(remaining)) + }) + + t.Run("accepted body is restored", func(t *testing.T) { + body := `{"invocation_id":"inv_body","status":"accepted"}` + original := &trackingReadCloser{Reader: strings.NewReader(body)} + resp := &http.Response{ + StatusCode: http.StatusAccepted, + Header: make(http.Header), + Body: original, + } + id, err := invocationIDFromResponse(resp) + require.NoError(t, err) + assert.Equal(t, "inv_body", id) + restored, err := io.ReadAll(resp.Body) + require.NoError(t, err) + assert.Equal(t, body, string(restored)) + assert.True(t, original.closed) + }) +} + +func TestInvocationsProtocolDispatchHTTP(t *testing.T) { + for _, operation := range []invocationOperation{invocationShow, invocationCancel, invocationFollow} { + t.Run(string(operation), func(t *testing.T) { + var methods []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + methods = append(methods, r.Method) + assert.Equal(t, "Bearer test-token", r.Header.Get("Authorization")) + assert.Equal(t, "v1", r.URL.Query().Get("api-version")) + assert.Empty(t, r.URL.Query().Get("agent_session_id")) + path := "/agents/agent/endpoint/protocols/invocations/inv_test" + if r.Method == http.MethodPost { + path += "/cancel" + } + assert.Equal(t, path, r.URL.Path) + _, _ = io.WriteString(w, `{"invocation_id":"inv_test","status":"completed"}`) + })) + defer server.Close() + action := &InvokeAction{flags: &invokeFlags{protocol: "invocations"}, credential: responseTestCredential{}} + rc := &remoteContext{projectEndpoint: server.URL, name: "agent", apiVersion: "v1"} + var output bytes.Buffer + err := action.runInvocationOperation(t.Context(), rc, "inv_test", operation, "json", &output) + switch operation { + case invocationShow: + require.NoError(t, err) + assert.JSONEq(t, `{"invocation_id":"inv_test","status":"completed"}`, output.String()) + assert.Equal(t, []string{"GET"}, methods) + case invocationCancel: + require.NoError(t, err) + assert.Contains(t, output.String(), "is completed") + assert.Equal(t, []string{"POST"}, methods) + case invocationFollow: + require.ErrorContains(t, err, "not supported") + assert.Empty(t, methods) + } + }) + } +} + +func TestInvocationLifecycleURLs(t *testing.T) { + assert.Equal( + t, + "https://example.test/agents/agent/endpoint/protocols/invocations/inv_1?api-version=v1", + buildInvocationLifecycleURL("https://example.test", "agent", "inv_1", "v1"), + ) + assert.Equal( + t, + "https://example.test/agents/agent/endpoint/protocols/invocations/inv_1/cancel?api-version=v1", + buildInvocationCancelURL("https://example.test", "agent", "inv_1", "v1"), + ) +} + +func TestPrintInvocationSnapshot(t *testing.T) { + var output bytes.Buffer + require.NoError(t, printInvocationSnapshot(&output, invocationSnapshotResult{ + snapshot: invocationSnapshot{ID: "inv_1", Status: "completed"}, + raw: []byte(`{"id":"inv_1","status":"completed"}`), + }, "json")) + assert.JSONEq(t, `{"id":"inv_1","status":"completed"}`, output.String()) + + output.Reset() + require.NoError(t, printInvocationSnapshot(&output, invocationSnapshotResult{ + snapshot: invocationSnapshot{ID: "inv_1", Status: "completed"}, + }, "table")) + assert.Contains(t, output.String(), "Invocation ID inv_1") + assert.Contains(t, output.String(), "Status completed") +} + +func TestInvocationStateStoreRoundTrip(t *testing.T) { + server := newInvokeUserConfigServer() + client := newInvokeTestAzdClient(t, server) + store := newInvocationStateStore(client) + want := savedInvocation{InvocationID: "inv_123"} + + require.NoError(t, store.Save(t.Context(), "agent-a", want)) + got, err := store.Get(t.Context(), "agent-a") + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, want, *got) + + require.NoError(t, store.Delete(t.Context(), "agent-a")) + got, err = store.Get(t.Context(), "agent-a") + require.NoError(t, err) + assert.Nil(t, got) +} diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_response_test.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_response_test.go index ec1bd71cd3a..90f6d8e2634 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_response_test.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_response_test.go @@ -713,6 +713,10 @@ func TestCleanupAgentStateForKey(t *testing.T) { agentKey: {ResponseID: "resp_123"}, otherKey: {ResponseID: "resp_other"}, }) + server.setJSON(t, invocationsConfigPath, map[string]savedInvocation{ + agentKey: {InvocationID: "inv_123"}, + otherKey: {InvocationID: "inv_other"}, + }) client := newInvokeTestAzdClient(t, server) require.True(t, cleanupAgentStateForKey(t.Context(), client, agentKey)) @@ -731,6 +735,11 @@ func TestCleanupAgentStateForKey(t *testing.T) { server.getJSON(t, responsesConfigPath, &responses) assert.NotContains(t, responses, agentKey) assert.Equal(t, "resp_other", responses[otherKey].ResponseID) + + var invocations map[string]savedInvocation + server.getJSON(t, invocationsConfigPath, &invocations) + assert.NotContains(t, invocations, agentKey) + assert.Equal(t, "inv_other", invocations[otherKey].InvocationID) } func TestCleanupPromptAgentStateUsesInvocationKey(t *testing.T) { diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/listen.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/listen.go index 24f5e5c1e89..bc271d2cddd 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/listen.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/listen.go @@ -571,7 +571,7 @@ func warnLegacySimpleTeamsArtifacts(proj *azdext.ProjectConfig, svc *azdext.Serv )) } -// postdownHandler cleans up saved session, conversation, and current Response state for agent services +// postdownHandler cleans up saved session, conversation, Response, and Invocation state for agent services // that were torn down. This is best-effort — failures are logged but do not block azd down. func postdownHandler(ctx context.Context, azdClient *azdext.AzdClient, args *azdext.ProjectEventArgs) error { envResp, err := azdClient.Environment().GetCurrent(ctx, &azdext.EmptyRequest{}) @@ -588,7 +588,7 @@ func postdownHandler(ctx context.Context, azdClient *azdext.AzdClient, args *azd } if cleanupAgentState(ctx, azdClient, envName, svc.Name) { - fmt.Printf("Cleaned up saved session, conversation, and current Response for agent %q\n", svc.Name) + fmt.Printf("Cleaned up saved session, conversation, Response, and Invocation state for agent %q\n", svc.Name) } } @@ -701,7 +701,7 @@ func cleanupPromptAgentState( return cleanupAgentStateForKey(ctx, azdClient, agentKey) } -// cleanupAgentState removes saved session, conversation, and current Response state for a +// cleanupAgentState removes saved session, conversation, Response, and Invocation state for a // single agent service. Returns true if cleanup succeeded, false otherwise. // Shared by postdownHandler and delete command. func cleanupAgentState(ctx context.Context, azdClient *azdext.AzdClient, envName, serviceName string) bool { @@ -733,6 +733,10 @@ func cleanupAgentStateForKey(ctx context.Context, azdClient *azdext.AzdClient, a log.Printf("cleanupAgentState: failed to clean current Response for %s: %v", agentKey, err) failed = true } + if err := newInvocationStateStore(azdClient).Delete(ctx, agentKey); err != nil { + log.Printf("cleanupAgentState: failed to clean current Invocation for %s: %v", agentKey, err) + failed = true + } return !failed } diff --git a/cli/azd/extensions/azure.ai.agents/internal/exterrors/codes.go b/cli/azd/extensions/azure.ai.agents/internal/exterrors/codes.go index 2f4c4f68dc2..64f50ed3523 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/exterrors/codes.go +++ b/cli/azd/extensions/azure.ai.agents/internal/exterrors/codes.go @@ -44,8 +44,11 @@ const ( CodeInvalidPositionalArg = "invalid_positional_arg" ) -// CodeInvalidResponseState identifies malformed locally saved Response state. -const CodeInvalidResponseState = "invalid_response_state" +// Error codes for malformed locally saved protocol resource state. +const ( + CodeInvalidResponseState = "invalid_response_state" + CodeInvalidInvocationState = "invalid_invocation_state" +) // CodeInvalidEnvironmentVariableName identifies a hosted-agent // environment variable name rejected by the service contract. @@ -223,12 +226,15 @@ const ( OpPublishTeamsApp = "publish_teams_app" ) -// Operation names for Responses protocol resources. +// Operation names for Responses and Invocations protocol resources. const ( - OpReadResponseState = "read_response_state" - OpShowResponse = "show_response" - OpFollowResponse = "follow_response" - OpCancelResponse = "cancel_response" + OpReadResponseState = "read_response_state" + OpShowResponse = "show_response" + OpFollowResponse = "follow_response" + OpCancelResponse = "cancel_response" + OpReadInvocationState = "read_invocation_state" + OpShowInvocation = "show_invocation" + OpCancelInvocation = "cancel_invocation" ) // Error codes for eval and optimize operations. From a66738f7dea733876518ade3fe84077eec5b0e91 Mon Sep 17 00:00:00 2001 From: Wei Meng Date: Thu, 10 Sep 2026 12:40:32 +0800 Subject: [PATCH 2/5] refactor(ai): document Invocation helpers and drop lifecycle custom headers --- .../internal/cmd/agent_endpoint.go | 2 ++ .../internal/cmd/invoke_invocation.go | 15 +++++++++++++-- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/agent_endpoint.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/agent_endpoint.go index 926a48603a0..75d6689877f 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/agent_endpoint.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/agent_endpoint.go @@ -198,6 +198,7 @@ func buildInvocationsURL(projectEndpoint, agentName, apiVersion, sid string) str return invURL } +// buildInvocationLifecycleURL builds the Invocation retrieval URL. func buildInvocationLifecycleURL( projectEndpoint string, agentName string, @@ -216,6 +217,7 @@ func buildInvocationLifecycleURL( ) } +// buildInvocationCancelURL builds the Invocation cancellation URL. func buildInvocationCancelURL(projectEndpoint, agentName, invocationID, apiVersion string) string { lifecycleURL := buildInvocationLifecycleURL(projectEndpoint, agentName, invocationID, apiVersion) parts := strings.SplitN(lifecycleURL, "?", 2) diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go index bee7baa0d65..80d09c5506b 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go @@ -21,6 +21,7 @@ import ( "github.com/azure/azure-dev/cli/azd/pkg/azdext" ) +// invocationIDFromResponse extracts an ID without consuming the accepted response body. func invocationIDFromResponse(resp *http.Response) (string, error) { if invocationID := resp.Header.Get("x-agent-invocation-id"); invocationID != "" { return invocationID, nil @@ -56,6 +57,7 @@ type invocationSnapshotResult struct { raw []byte } +// runInvocationsProtocolOperation executes Invocations-protocol lifecycle requests. func (a *InvokeAction) runInvocationsProtocolOperation( ctx context.Context, rc *remoteContext, @@ -82,6 +84,7 @@ func (a *InvokeAction) runInvocationsProtocolOperation( } } +// classifyInvocationStateReadError adds guidance to saved-state failures. func classifyInvocationStateReadError(cause error) error { if _, ok := errors.AsType[*azdext.ConfigError](cause); !ok { return exterrors.FromHost(cause, exterrors.OpReadInvocationState, "reading current Invocation state failed") @@ -93,6 +96,7 @@ func classifyInvocationStateReadError(cause error) error { ) } +// cancelInvocation requests cancellation and confirms terminal state on rejection. func (a *InvokeAction) cancelInvocation( ctx context.Context, rc *remoteContext, invocationID string, writer io.Writer, ) error { @@ -106,7 +110,6 @@ func (a *InvokeAction) cancelInvocation( return fmt.Errorf("create Invocation cancel request: %w", err) } req.Header.Set("Authorization", "Bearer "+token) - applyCustomHeaders(req, a.clientHeaders) applyRemoteUserIdentityHeader(req, &a.flags.userIdentityFlags) //nolint:gosec // URL is built from a validated Foundry endpoint. @@ -148,6 +151,7 @@ func (a *InvokeAction) cancelInvocation( return err } +// getInvocation retrieves an Invocation and validates its identity. func (a *InvokeAction) getInvocation( ctx context.Context, rc *remoteContext, @@ -163,7 +167,6 @@ func (a *InvokeAction) getInvocation( return invocationSnapshotResult{}, fmt.Errorf("create Invocation show request: %w", err) } req.Header.Set("Authorization", "Bearer "+token) - applyCustomHeaders(req, a.clientHeaders) applyRemoteUserIdentityHeader(req, &a.flags.userIdentityFlags) //nolint:gosec // URL is built from a validated Foundry endpoint. @@ -211,10 +214,12 @@ type invocationLifecycleHTTPError struct { body []byte } +// Error describes the failed lifecycle HTTP request. func (e *invocationLifecycleHTTPError) Error() string { return fmt.Sprintf("%s %s failed with HTTP %d: %s\n%s", e.method, e.requestURL, e.statusCode, e.status, e.body) } +// classifyInvocationLifecycleError translates HTTP failures into structured service errors. func classifyInvocationLifecycleError(cause error, operation, label string) error { if cause == nil { return nil @@ -238,6 +243,7 @@ func classifyInvocationLifecycleError(cause error, operation, label string) erro return serviceErr } +// printInvocationSnapshot writes an Invocation as JSON or a table. func printInvocationSnapshot(writer io.Writer, result invocationSnapshotResult, format string) error { if format != "table" { var formatted any @@ -264,6 +270,7 @@ func printInvocationSnapshot(writer io.Writer, result invocationSnapshotResult, return table.Flush() } +// isTerminalInvocationStatus reports whether an Invocation has finished. func isTerminalInvocationStatus(status string) bool { switch status { case "completed", "failed", "cancelled", "canceled": @@ -283,10 +290,12 @@ type invocationStateStore struct { client *azdext.AzdClient } +// newInvocationStateStore creates the current-Invocation ID store. func newInvocationStateStore(client *azdext.AzdClient) *invocationStateStore { return &invocationStateStore{client: client} } +// Get returns the agent's current Invocation, or nil if none is saved. func (s *invocationStateStore) Get(ctx context.Context, agentKey string) (*savedInvocation, error) { config, err := azdext.NewConfigHelper(s.client) if err != nil { @@ -307,6 +316,7 @@ func (s *invocationStateStore) Get(ctx context.Context, agentKey string) (*saved return &record, nil } +// Save replaces the agent's current Invocation ID. func (s *invocationStateStore) Save(ctx context.Context, agentKey string, record savedInvocation) error { config, err := azdext.NewConfigHelper(s.client) if err != nil { @@ -327,6 +337,7 @@ func (s *invocationStateStore) Save(ctx context.Context, agentKey string, record return nil } +// Delete removes the agent's current Invocation selection. func (s *invocationStateStore) Delete(ctx context.Context, agentKey string) error { config, err := azdext.NewConfigHelper(s.client) if err != nil { From e41aaa2fff54b1976d52aba23338f0b95d81a75f Mon Sep 17 00:00:00 2001 From: Wei Meng Date: Thu, 10 Sep 2026 14:26:54 +0800 Subject: [PATCH 3/5] docs(ai): clarify Invocation behavior and scenario mapping --- .../azure.ai.agents/docs/specs/long-running-agent-invoke.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cli/azd/extensions/azure.ai.agents/docs/specs/long-running-agent-invoke.md b/cli/azd/extensions/azure.ai.agents/docs/specs/long-running-agent-invoke.md index 6ec97dbb59b..67e796a8ffc 100644 --- a/cli/azd/extensions/azure.ai.agents/docs/specs/long-running-agent-invoke.md +++ b/cli/azd/extensions/azure.ai.agents/docs/specs/long-running-agent-invoke.md @@ -42,7 +42,7 @@ azd ai agent invocations show --protocol invocations --id azd ai agent invocations cancel --protocol invocations --id ``` -The Invocations-protocol lifecycle implementation is included in PR #9901. Its existing synchronous, SSE, raw, and `202 Accepted` polling behavior on create is unchanged. Creation captures the ID from `x-agent-invocation-id`, or from the `invocation_id` property of a `202 Accepted` body. The body is preserved for the existing polling and raw-output handlers. Show and cancel use the resolved endpoint/API version, without inheriting session context from the create. +The Invocations-protocol lifecycle commands preserve the existing synchronous, SSE, raw, and `202 Accepted` polling behavior on create. Creation captures the ID from `x-agent-invocation-id`, or from the `invocation_id` property of a `202 Accepted` body. The body is preserved for the existing polling and raw-output handlers. Show and cancel use the resolved endpoint/API version, without inheriting session context from the create. ### Protocol selection From 730d57055e86993cec672e6bf4438bf6917b9ddb Mon Sep 17 00:00:00 2001 From: Wei Meng Date: Fri, 11 Sep 2026 09:26:32 +0800 Subject: [PATCH 4/5] fix(ai): validate Invocation headers and explain unsupported cancellation --- .../internal/cmd/invoke_invocation.go | 55 +++++++++--- .../internal/cmd/invoke_invocation_test.go | 89 +++++++++++++++++++ 2 files changed, 133 insertions(+), 11 deletions(-) diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go index 80d09c5506b..f4878616eaf 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation.go @@ -13,6 +13,7 @@ import ( "net/http" "net/url" "strconv" + "strings" "text/tabwriter" "time" @@ -188,21 +189,39 @@ func (a *InvokeAction) getInvocation( body: body, } } - var snapshot invocationSnapshot - if err := json.Unmarshal(body, &snapshot); err != nil { - return invocationSnapshotResult{}, fmt.Errorf("decode Invocation: %w", err) - } - actualID := snapshot.ID - if actualID == "" { - actualID = snapshot.InvocationID - } + actualID := resp.Header.Get("x-agent-invocation-id") if actualID != "" && actualID != invocationID { return invocationSnapshotResult{}, fmt.Errorf( - "Invocation ID %q does not match requested ID %q", - actualID, - invocationID, + "Invocation ID %q does not match requested ID %q", actualID, invocationID, ) } + var snapshot invocationSnapshot + if actualID != "" { + // Payload IDs belong to the handler; the protocol header takes precedence. + var result struct { + Status string `json:"status"` + } + if err := json.Unmarshal(body, &result); err != nil { + return invocationSnapshotResult{}, fmt.Errorf("decode Invocation: %w", err) + } + snapshot = invocationSnapshot{InvocationID: actualID, Status: result.Status} + } else { + if err := json.Unmarshal(body, &snapshot); err != nil { + return invocationSnapshotResult{}, fmt.Errorf("decode Invocation: %w", err) + } + actualID = snapshot.ID + if actualID == "" { + actualID = snapshot.InvocationID + } + if actualID != "" && actualID != invocationID { + return invocationSnapshotResult{}, fmt.Errorf( + "Invocation ID %q does not match requested ID %q", actualID, invocationID, + ) + } + if actualID == "" { + snapshot.InvocationID = invocationID + } + } return invocationSnapshotResult{snapshot: snapshot, raw: body}, nil } @@ -240,6 +259,20 @@ func classifyInvocationLifecycleError(cause error, operation, label string) erro "", ) serviceErr.StatusCode = httpErr.statusCode + // Recognize the known agent-server rejection without exposing arbitrary handler content. + if operation == exterrors.OpCancelInvocation && httpErr.statusCode == http.StatusNotFound && + len(httpErr.body) <= 16*1024 { + var body struct { + Error struct { + Code string `json:"code"` + Message string `json:"message"` + } `json:"error"` + } + if json.Unmarshal(httpErr.body, &body) == nil && body.Error.Code == "not_found" && + strings.TrimSpace(body.Error.Message) == "cancel_invocation not implemented" { + serviceErr.Message = "This agent does not support cancelling invocations." + } + } return serviceErr } diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go index 696af8f5f64..2725968cc68 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go @@ -5,12 +5,16 @@ package cmd import ( "bytes" + "errors" "io" "net/http" "net/http/httptest" "strings" "testing" + "azureaiagent/internal/exterrors" + + "github.com/azure/azure-dev/cli/azd/pkg/azdext" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -102,6 +106,91 @@ func TestInvocationsProtocolDispatchHTTP(t *testing.T) { } } +func TestInvocationSnapshotIdentityHTTP(t *testing.T) { + for _, tt := range []struct { + name, header, body, wantErr string + }{ + {name: "header only", header: "inv_test", body: `{"status":"completed","result":"done"}`}, + {name: "header overrides handler IDs", header: "inv_test", + body: `{"id":123,"invocation_id":"handler-owned-id","status":"completed"}`}, + {name: "header mismatch even with matching body", header: "inv_other", + body: `{"invocation_id":"inv_test","status":"completed"}`, wantErr: "does not match requested ID"}, + {name: "body invocation ID fallback", body: `{"invocation_id":"inv_test","status":"completed"}`}, + {name: "body ID fallback", body: `{"id":"inv_test","status":"completed"}`}, + {name: "body mismatch", body: `{"invocation_id":"inv_other"}`, wantErr: "does not match requested ID"}, + {name: "requested ID when no identity supplied", body: `{"status":"completed","result":"done"}`}, + } { + t.Run(tt.name, func(t *testing.T) { + calls := 0 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + assert.Equal(t, http.MethodGet, r.Method) + assert.Equal(t, "/agents/agent/endpoint/protocols/invocations/inv_test", r.URL.Path) + if tt.header != "" { + w.Header().Set("x-agent-invocation-id", tt.header) + } + _, _ = io.WriteString(w, tt.body) + })) + defer server.Close() + action := &InvokeAction{flags: &invokeFlags{protocol: "invocations"}, credential: responseTestCredential{}} + rc := &remoteContext{projectEndpoint: server.URL, name: "agent", apiVersion: "v1"} + result, err := action.getInvocation(t.Context(), rc, "inv_test") + assert.Equal(t, 1, calls) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return + } + require.NoError(t, err) + assert.Equal(t, tt.body, string(result.raw), "do not rewrite the handler's JSON") + var output bytes.Buffer + require.NoError(t, printInvocationSnapshot(&output, result, "table")) + assert.Contains(t, output.String(), "Invocation ID inv_test") + output.Reset() + require.NoError(t, printInvocationSnapshot(&output, result, "json")) + assert.JSONEq(t, tt.body, output.String()) + }) + } +} + +func TestInvocationCancelUnsupportedHTTP(t *testing.T) { + var methods []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + methods = append(methods, r.Method) + if r.Method == http.MethodPost { + w.WriteHeader(http.StatusNotFound) + _, _ = io.WriteString(w, `{"error":{"code":"not_found","message":"cancel_invocation not implemented"}}`) + return + } + w.Header().Set("x-agent-invocation-id", "inv_test") + _, _ = io.WriteString(w, `{"status":"running"}`) + })) + defer server.Close() + action := &InvokeAction{flags: &invokeFlags{protocol: "invocations"}, credential: responseTestCredential{}} + rc := &remoteContext{projectEndpoint: server.URL, name: "agent", apiVersion: "v1"} + err := action.runInvocationOperation(t.Context(), rc, "inv_test", invocationCancel, "", io.Discard) + require.EqualError(t, err, "This agent does not support cancelling invocations.") + serviceErr, ok := errors.AsType[*azdext.ServiceError](err) + require.True(t, ok) + assert.Equal(t, http.StatusNotFound, serviceErr.StatusCode) + assert.NotEmpty(t, serviceErr.ServiceName) + assert.Equal(t, []string{"POST", "GET"}, methods) +} + +func TestInvocationErrorDetailRemainsBounded(t *testing.T) { + for _, body := range []string{ + `{"error":{"code":"not_found","message":"private handler content"}}`, + `not JSON`, + `{"error":{"code":"not_found","message":"cancel_invocation not implemented"},"padding":"` + + strings.Repeat("x", 16*1024) + `"}`, + } { + err := classifyInvocationLifecycleError(&invocationLifecycleHTTPError{ + method: http.MethodPost, requestURL: "https://example.test/invocations/inv_test/cancel", + statusCode: http.StatusNotFound, status: "404 Not Found", body: []byte(body), + }, exterrors.OpCancelInvocation, "cancelling Invocation") + require.EqualError(t, err, "cancelling Invocation failed with HTTP 404: 404 Not Found") + } +} + func TestInvocationLifecycleURLs(t *testing.T) { assert.Equal( t, From 202258b0100622cf199b74f903d10d8e13bd10e1 Mon Sep 17 00:00:00 2001 From: Wei Meng Date: Mon, 14 Sep 2026 10:40:19 +0800 Subject: [PATCH 5/5] test(ai): cover Invocation create persistence and terminal cancellation --- .../internal/cmd/invoke_invocation_test.go | 173 ++++++++++++++++++ 1 file changed, 173 insertions(+) diff --git a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go index 2725968cc68..9021db7d45f 100644 --- a/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go +++ b/cli/azd/extensions/azure.ai.agents/internal/cmd/invoke_invocation_test.go @@ -68,6 +68,116 @@ func TestInvocationIDFromResponse(t *testing.T) { }) } +func TestInvocationsRemotePersistsCurrentID(t *testing.T) { + const agentKey = "target-agent" + const invocationID = "inv_created" + for _, tt := range []struct { + name string + status int + headerID string + contentType string + body string + wantID string + wantOutput string + }{ + { + name: "sync header ID", status: http.StatusOK, headerID: invocationID, + body: `{"result":"sync-result"}`, wantID: invocationID, wantOutput: "sync-result", + }, + { + name: "stream header ID", status: http.StatusOK, headerID: invocationID, + contentType: "text/event-stream", body: "data: stream-result\n\ndata: [DONE]\n\n", + wantID: invocationID, wantOutput: "stream-result", + }, + { + name: "accepted body ID", status: http.StatusAccepted, + body: `{"invocation_id":"inv_created","status":"running","result":"accepted-result"}`, + wantID: invocationID, wantOutput: "terminal-result", + }, + { + name: "HTTP failure preserves selection", status: http.StatusInternalServerError, headerID: invocationID, + body: `{"error":{"message":"request rejected"}}`, wantID: "inv_previous", + }, + { + name: "missing ID preserves selection", status: http.StatusOK, + body: `{"result":"sync-result"}`, wantID: "inv_previous", wantOutput: "sync-result", + }, + } { + for _, format := range []string{"default", outputRaw} { + t.Run(tt.name+"/"+format, func(t *testing.T) { + config := newInvokeUserConfigServer() + config.setJSON(t, invocationsConfigPath, map[string]savedInvocation{ + agentKey: {InvocationID: "inv_previous"}, "other-agent": {InvocationID: "inv_other"}, + }) + config.setJSON(t, responsesConfigPath, map[string]savedResponse{ + agentKey: {ResponseID: "resp_previous"}, + }) + client := newInvokeTestAzdClient(t, config) + var methods []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + methods = append(methods, r.Method) + assert.Equal(t, "Bearer test-token", r.Header.Get("Authorization")) + assert.Equal(t, "v1", r.URL.Query().Get("api-version")) + path := "/agents/agent/endpoint/protocols/invocations" + if r.Method == http.MethodGet { + assert.Equal(t, http.StatusAccepted, tt.status) + assert.Equal(t, path+"/"+invocationID, r.URL.Path) + _, _ = io.WriteString(w, + `{"invocation_id":"inv_created","status":"completed","result":"terminal-result"}`) + return + } + assert.Equal(t, http.MethodPost, r.Method) + assert.Equal(t, path, r.URL.Path) + if tt.headerID != "" { + w.Header().Set("x-agent-invocation-id", tt.headerID) + } + if tt.contentType != "" { + w.Header().Set("Content-Type", tt.contentType) + } + w.WriteHeader(tt.status) + _, _ = io.WriteString(w, tt.body) + })) + defer server.Close() + action := &InvokeAction{ + flags: &invokeFlags{message: `{"input":"test"}`, protocol: "invocations", outputFmt: format}, + credential: responseTestCredential{}, + // Endpoint mode avoids project-only OpenAPI caching in this HTTP/config integration test. + endpoint: &parsedAgentEndpoint{}, + resolvedRemoteContext: &remoteContext{ + projectEndpoint: server.URL, name: "agent", serviceName: "service", apiVersion: "v1", + agentKey: agentKey, azdClient: client, + }, + } + var invokeErr error + output := withCapturedStdout(t, func() { invokeErr = action.invocationsRemote(t.Context()) }) + if tt.status >= http.StatusBadRequest { + require.ErrorContains(t, invokeErr, "HTTP 500") + } else { + require.NoError(t, invokeErr) + assert.Contains(t, output, tt.wantOutput) + } + wantMethods := []string{"POST"} + if tt.status == http.StatusAccepted { + wantMethods = append(wantMethods, "GET") + } + assert.Equal(t, wantMethods, methods) + if format == outputRaw { + assert.Contains(t, output, tt.body, "ID extraction must preserve the original response body") + assert.NotContains(t, output, "Invocation: ") + } + var saved map[string]savedInvocation + config.getJSON(t, invocationsConfigPath, &saved) + assert.Equal(t, map[string]savedInvocation{ + agentKey: {InvocationID: tt.wantID}, "other-agent": {InvocationID: "inv_other"}, + }, saved) + var responses map[string]savedResponse + config.getJSON(t, responsesConfigPath, &responses) + assert.Equal(t, map[string]savedResponse{agentKey: {ResponseID: "resp_previous"}}, responses) + }) + } + } +} + func TestInvocationsProtocolDispatchHTTP(t *testing.T) { for _, operation := range []invocationOperation{invocationShow, invocationCancel, invocationFollow} { t.Run(string(operation), func(t *testing.T) { @@ -176,6 +286,69 @@ func TestInvocationCancelUnsupportedHTTP(t *testing.T) { assert.Equal(t, []string{"POST", "GET"}, methods) } +func TestInvocationCancelTerminalFallbackHTTP(t *testing.T) { + for _, tt := range []struct { + name string + status string + getStatus int + headerID string + wantNoWork bool + }{ + {name: "completed", status: "completed", wantNoWork: true}, + {name: "failed", status: "failed", wantNoWork: true}, + {name: "cancelled", status: "cancelled", wantNoWork: true}, + {name: "canceled alias", status: "canceled", wantNoWork: true}, + {name: "still active", status: "running"}, + {name: "GET failed", getStatus: http.StatusNotFound}, + {name: "terminal response for wrong ID", status: "completed", headerID: "inv_other"}, + } { + t.Run(tt.name, func(t *testing.T) { + var methods []string + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + methods = append(methods, r.Method) + assert.Equal(t, "Bearer test-token", r.Header.Get("Authorization")) + assert.Equal(t, "v1", r.URL.Query().Get("api-version")) + assert.Empty(t, r.URL.Query().Get("agent_session_id")) + path := "/agents/agent/endpoint/protocols/invocations/inv_test" + if r.Method == http.MethodPost { + assert.Equal(t, path+"/cancel", r.URL.Path) + w.WriteHeader(http.StatusConflict) + _, _ = io.WriteString(w, `{"error":{"message":"cancellation rejected"}}`) + return + } + assert.Equal(t, http.MethodGet, r.Method) + assert.Equal(t, path, r.URL.Path) + if tt.getStatus != 0 { + w.WriteHeader(tt.getStatus) + return + } + headerID := tt.headerID + if headerID == "" { + headerID = "inv_test" + } + w.Header().Set("x-agent-invocation-id", headerID) + _, _ = io.WriteString(w, `{"status":"`+tt.status+`"}`) + })) + defer server.Close() + action := &InvokeAction{flags: &invokeFlags{protocol: "invocations"}, credential: responseTestCredential{}} + rc := &remoteContext{projectEndpoint: server.URL, name: "agent", apiVersion: "v1"} + var output bytes.Buffer + err := action.runInvocationOperation(t.Context(), rc, "inv_test", invocationCancel, "", &output) + assert.Equal(t, []string{"POST", "GET"}, methods) + if tt.wantNoWork { + require.NoError(t, err) + assert.Equal(t, "Invocation inv_test is already "+tt.status+"; nothing to cancel.\n", output.String()) + } else { + require.Error(t, err) + serviceErr, ok := errors.AsType[*azdext.ServiceError](err) + require.True(t, ok) + assert.Equal(t, http.StatusConflict, serviceErr.StatusCode, "preserve the original cancel rejection") + assert.Empty(t, output.String()) + } + }) + } +} + func TestInvocationErrorDetailRemainsBounded(t *testing.T) { for _, body := range []string{ `{"error":{"code":"not_found","message":"private handler content"}}`,