Refactor session ID handling in response capture and invocation methods. Fixes #7573.

Author: therealjohn

cli/azd/extensions/azure.ai.agents/internal/cmd/helpers.go

@@ -199,16 +199,12 @@ func contextMap(agentCtx *AgentLocalContext, field string) map[string]string {
 // printSessionStatus prints the session line for the invoke banner.
 // label is the formatted prefix (e.g. "Session:  " or "Session:      ").
 func printSessionStatus(label, sid string) {
-	if sid != "" {
-		fmt.Printf("%s%s\n", label, sid)
-	} else {
-		fmt.Printf("%s(new — server will assign)\n", label)
-	}
+	fmt.Printf("%s%s\n", label, sid)
 }
 
 // captureResponseSession reads the x-agent-session-id header from a response
-// and saves it when the caller had no pre-existing session (sid == "").
-// label is the formatted prefix for printing (e.g. "Session:  ").
+// and updates the persisted session when the server returns a different ID
+// than what the client sent (or assigns one when none was sent).
 func captureResponseSession(
 	ctx context.Context,
 	azdClient *azdext.AzdClient,
@@ -217,12 +213,19 @@ func captureResponseSession(
 	resp *http.Response,
 	label string,
 ) {
-	if sid != "" || azdClient == nil {
+	if azdClient == nil {
 		return
 	}
-	if newSid := resp.Header.Get("x-agent-session-id"); newSid != "" {
-		saveContextValue(ctx, azdClient, agentName, newSid, "sessions")
+	newSid := resp.Header.Get("x-agent-session-id")
+	if newSid == "" || newSid == sid {
+		return
+	}
+	// Server assigned or reassigned a session ID — update persisted value.
+	saveContextValue(ctx, azdClient, agentName, newSid, "sessions")
+	if sid == "" {
 		fmt.Printf("%s%s (assigned by server)\n", label, newSid)
+	} else {
+		fmt.Printf("%s%s (reassigned by server)\n", label, newSid)
 	}
 }
 

cli/azd/extensions/azure.ai.agents/internal/cmd/helpers_test.go

@@ -4,6 +4,7 @@
 package cmd
 
 import (
+	"net/http"
 	"os"
 	"path/filepath"
 	"testing"
@@ -178,3 +179,131 @@ func TestToServiceKey(t *testing.T) {
 		})
 	}
 }
+
+func TestCaptureResponseSession_NilClient(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name      string
+		sid       string
+		headerVal string
+	}{
+		{name: "no header", sid: "", headerVal: ""},
+		{name: "header present but nil client", sid: "", headerVal: "server-session-abc"},
+		{name: "client sid set with header", sid: "existing-session", headerVal: "server-session-abc"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			resp := &http.Response{Header: http.Header{}}
+			if tt.headerVal != "" {
+				resp.Header.Set("x-agent-session-id", tt.headerVal)
+			}
+
+			// Must not panic with nil azdClient.
+			captureResponseSession(t.Context(), nil, "test-agent", tt.sid, resp, "Session: ")
+		})
+	}
+}
+
+func TestLoadSaveLocalContext(t *testing.T) {
+	t.Parallel()
+
+	t.Run("round trip", func(t *testing.T) {
+		t.Parallel()
+
+		dir := t.TempDir()
+		configPath := filepath.Join(dir, ConfigFile)
+
+		agentCtx := &AgentLocalContext{
+			AgentName: "my-agent",
+			Sessions:  map[string]string{"agent1": "sess-123"},
+		}
+
+		if err := saveLocalContext(agentCtx, configPath); err != nil {
+			t.Fatalf("saveLocalContext failed: %v", err)
+		}
+
+		loaded := loadLocalContext(configPath)
+		if loaded.AgentName != "my-agent" {
+			t.Errorf("AgentName = %q, want %q", loaded.AgentName, "my-agent")
+		}
+		if loaded.Sessions["agent1"] != "sess-123" {
+			t.Errorf("Sessions[agent1] = %q, want %q", loaded.Sessions["agent1"], "sess-123")
+		}
+	})
+
+	t.Run("missing file returns empty context", func(t *testing.T) {
+		t.Parallel()
+
+		loaded := loadLocalContext(filepath.Join(t.TempDir(), "nonexistent.json"))
+		if loaded.Sessions != nil {
+			t.Errorf("expected nil Sessions for missing file, got %v", loaded.Sessions)
+		}
+	})
+
+	t.Run("corrupt file returns empty context", func(t *testing.T) {
+		t.Parallel()
+
+		dir := t.TempDir()
+		configPath := filepath.Join(dir, ConfigFile)
+		if err := os.WriteFile(configPath, []byte("{bad json"), 0600); err != nil {
+			t.Fatalf("failed to write corrupt file: %v", err)
+		}
+
+		loaded := loadLocalContext(configPath)
+		if loaded.Sessions != nil {
+			t.Errorf("expected nil Sessions for corrupt file, got %v", loaded.Sessions)
+		}
+	})
+}
+
+func TestContextMap(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name  string
+		field string
+	}{
+		{name: "sessions", field: "sessions"},
+		{name: "conversations", field: "conversations"},
+		{name: "invocations", field: "invocations"},
+		{name: "unknown", field: "unknown"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			agentCtx := &AgentLocalContext{}
+			m := contextMap(agentCtx, tt.field)
+			if m == nil {
+				t.Fatal("expected non-nil map")
+			}
+			m["key"] = "value"
+
+			// For known fields, the map should be stored on the struct.
+			switch tt.field {
+			case "sessions":
+				if agentCtx.Sessions["key"] != "value" {
+					t.Error("sessions map not stored on struct")
+				}
+			case "conversations":
+				if agentCtx.Conversations["key"] != "value" {
+					t.Error("conversations map not stored on struct")
+				}
+			case "invocations":
+				if agentCtx.Invocations["key"] != "value" {
+					t.Error("invocations map not stored on struct")
+				}
+			case "unknown":
+				// Detached map — verify it doesn't affect any struct field.
+				if agentCtx.Sessions != nil || agentCtx.Conversations != nil || agentCtx.Invocations != nil {
+					t.Error("unknown field should not initialize struct maps")
+				}
+			}
+		})
+	}
+}

cli/azd/extensions/azure.ai.agents/internal/cmd/invoke.go

@@ -356,8 +356,8 @@ func (a *InvokeAction) responsesRemote(ctx context.Context) error {
 	}
 
 	// Session ID — routes to the same microVM container instance.
-	// When empty, let the server assign one.
-	sid, err := resolveStoredID(ctx, azdClient, name, a.flags.session, a.flags.newSession, "sessions", false)
+	// Generated client-side (UUID) so it is always persisted locally, even if the request fails.
+	sid, err := resolveStoredID(ctx, azdClient, name, a.flags.session, a.flags.newSession, "sessions", true)
 	if err != nil {
 		return err
 	}
@@ -542,8 +542,9 @@ func (a *InvokeAction) invocationsRemote(ctx context.Context) error {
 		return err
 	}
 
-	// Session ID — routes to the same container instance
-	sid, err := resolveStoredID(ctx, azdClient, name, a.flags.session, a.flags.newSession, "sessions", false)
+	// Session ID — routes to the same container instance.
+	// Generated client-side (UUID) so it is always persisted locally, even if the request fails.
+	sid, err := resolveStoredID(ctx, azdClient, name, a.flags.session, a.flags.newSession, "sessions", true)
 	if err != nil {
 		return err
 	}