shelley/ui: fix context window bar not updating on conversation switch

Philip Zeyliger created 4 months ago

Prompt: When you navigate between conversations, the "progress bar" for the
context window seems to get lost or reset to 0. If I navigate between
conversations, it should show me the progress bar as of the conversation I'm on!

When navigating between conversations, the context window progress bar
would sometimes not update correctly. This was because the server uses
'omitempty' on the context_window_size field, so when it's 0 (e.g., for
conversations with no agent responses yet), the field is omitted from
the JSON response.

Previously, the frontend checked 'typeof response.context_window_size === "number"'
which would be false when the field was omitted, leaving the old value
in state. Now we always update the context window size when loading a
conversation, defaulting to 0 if the field is not present.

Added a test to verify context_window_size is correctly returned when
navigating between different conversations.

Change summary

server/context_window_navigation_test.go | 84 ++++++++++++++++++++++++++
server/testharness_test.go               |  1 
ui/src/components/ChatInterface.tsx      |  6 
3 files changed, 88 insertions(+), 3 deletions(-)

Detailed changes

server/context_window_navigation_test.go

@@ -0,0 +1,84 @@
+package server
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+// TestContextWindowSizePreservedOnNavigation tests that context window size
+// is correctly returned when loading different conversations
+func TestContextWindowSizePreservedOnNavigation(t *testing.T) {
+	h := NewTestHarness(t)
+	defer h.Close()
+
+	// Create first conversation and get a response
+	h.NewConversation("echo: first message", "/tmp")
+	resp1 := h.WaitResponse()
+	t.Logf("First conversation first response: %q", resp1)
+
+	// Get context window size for first conversation
+	firstConvID := h.convID
+	firstConvSize := h.GetContextWindowSize()
+	t.Logf("First conversation context window size: %d", firstConvSize)
+	if firstConvSize == 0 {
+		t.Fatal("expected non-zero context window size for first conversation")
+	}
+
+	// Create second conversation and get a response
+	h.NewConversation("echo: second message with much more text to ensure different context size", "/tmp")
+	resp2 := h.WaitResponse()
+	t.Logf("Second conversation first response: %q", resp2)
+
+	secondConvID := h.convID
+	secondConvSize := h.GetContextWindowSize()
+	t.Logf("Second conversation context window size: %d", secondConvSize)
+	if secondConvSize == 0 {
+		t.Fatal("expected non-zero context window size for second conversation")
+	}
+
+	// Now simulate "navigating" back to the first conversation by fetching it via GET
+	// This is what the UI does when switching conversations
+	req := httptest.NewRequest("GET", "/api/conversation/"+firstConvID, nil)
+	w := httptest.NewRecorder()
+	h.server.handleGetConversation(w, req, firstConvID)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("GET first conversation returned %d: %s", w.Code, w.Body.String())
+	}
+
+	var resp StreamResponse
+	if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
+		t.Fatalf("Failed to decode response: %v", err)
+	}
+
+	t.Logf("First conversation on navigation: context_window_size=%d, messages=%d",
+		resp.ContextWindowSize, len(resp.Messages))
+
+	if resp.ContextWindowSize != firstConvSize {
+		t.Errorf("context_window_size mismatch on navigation: got %d, want %d",
+			resp.ContextWindowSize, firstConvSize)
+	}
+
+	// Now navigate to second conversation
+	req = httptest.NewRequest("GET", "/api/conversation/"+secondConvID, nil)
+	w = httptest.NewRecorder()
+	h.server.handleGetConversation(w, req, secondConvID)
+
+	if w.Code != http.StatusOK {
+		t.Fatalf("GET second conversation returned %d: %s", w.Code, w.Body.String())
+	}
+
+	if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
+		t.Fatalf("Failed to decode response: %v", err)
+	}
+
+	t.Logf("Second conversation on navigation: context_window_size=%d, messages=%d",
+		resp.ContextWindowSize, len(resp.Messages))
+
+	if resp.ContextWindowSize != secondConvSize {
+		t.Errorf("context_window_size mismatch on navigation: got %d, want %d",
+			resp.ContextWindowSize, secondConvSize)
+	}
+}

server/testharness_test.go

@@ -85,6 +85,7 @@ func (h *TestHarness) NewConversation(msg, cwd string) *TestHarness {
 		h.t.Fatalf("NewConversation: failed to parse response: %v", err)
 	}
 	h.convID = resp.ConversationID
+	h.responsesCount = 0 // Reset for new conversation
 	return h
 }
 

ui/src/components/ChatInterface.tsx

@@ -522,9 +522,9 @@ function ChatInterface({
       const response = await api.getConversation(conversationId);
       setMessages(response.messages ?? []);
       setAgentWorking(Boolean(response.agent_working));
-      if (typeof response.context_window_size === "number") {
-        setContextWindowSize(response.context_window_size);
-      }
+      // Always update context window size when loading a conversation.
+      // If omitted from response (due to omitempty when 0), default to 0.
+      setContextWindowSize(response.context_window_size ?? 0);
       if (onConversationUpdate) {
         onConversationUpdate(response.conversation);
       }