fix(chatd): address round 4 review — double injection, test coverage, nits

kylecarbs · kylecarbs · commit a132d45e7509 · 2026-04-15T19:14:45.000Z
diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go
@@ -5328,6 +5328,7 @@ func (p *Server) runChat(
 			}
 			if reloadedInstruction != "" {
 				reloadedPrompt = chatprompt.InsertSystem(reloadedPrompt, reloadedInstruction)
+				instructionInjected = true
 			}
 			reloadedPrompt = renderPlanPathPrompt(reloadedPrompt, resolvePlanPathBlock(reloadCtx))
 			reloadedSkills := skillsFromParts(reloadedMsgs)
@@ -5350,20 +5351,20 @@ func (p *Server) runChat(
 			}
 			return reloadedPrompt, nil
 		},
-			DisableChainMode: func() {
-				chainModeActive = false
-			},
-			PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message {
-				if instructionInjected || instruction == "" {
-					return nil
-				}
-				instructionInjected = true
-				result := chatprompt.InsertSystem(msgs, instruction)
-				if skillIndex := chattool.FormatSkillIndex(skills); skillIndex != "" {
-					result = chatprompt.InsertSystem(result, skillIndex)
-				}
-				return result
-			},
+		DisableChainMode: func() {
+			chainModeActive = false
+		},
+		PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message {
+			if instructionInjected || instruction == "" {
+				return nil
+			}
+			instructionInjected = true
+			result := chatprompt.InsertSystem(msgs, instruction)
+			if skillIndex := chattool.FormatSkillIndex(skills); skillIndex != "" {
+				result = chatprompt.InsertSystem(result, skillIndex)
+			}
+			return result
+		},
 		OnRetry: func(
 			attempt int,
 			retryErr error,
diff --git a/coderd/x/chatd/chatloop/chatloop.go b/coderd/x/chatd/chatloop/chatloop.go
@@ -355,16 +355,16 @@ func Run(ctx context.Context, opts RunOptions) error {
 			// copy copies Message structs by value, so field
 			// reassignments in addAnthropicPromptCaching only
 			// affect the prepared slice.
-				if opts.PrepareMessages != nil {
-					if updated := opts.PrepareMessages(messages); updated != nil {
-						messages = updated
-					}
-				}
-				prepared := make([]fantasy.Message, len(messages))
-				copy(prepared, messages)
-				if applyAnthropicCaching {
-					addAnthropicPromptCaching(prepared)
+			if opts.PrepareMessages != nil {
+				if updated := opts.PrepareMessages(messages); updated != nil {
+					messages = updated
 				}
+			}
+			prepared := make([]fantasy.Message, len(messages))
+			copy(prepared, messages)
+			if applyAnthropicCaching {
+				addAnthropicPromptCaching(prepared)
+			}
 
 			call := fantasy.Call{
 				Prompt:           prepared,
diff --git a/coderd/x/chatd/chatloop/chatloop_test.go b/coderd/x/chatd/chatloop/chatloop_test.go
@@ -1601,7 +1601,7 @@ func TestRun_PrepareMessagesInjectsSystemContextMidLoop(t *testing.T) {
 
 	// Simulate: after the tool executes (step 0), instruction
 	// becomes available. PrepareMessages injects it before step 1.
-	var instructionInjected atomic.Bool
+	instructionInjected := make(chan struct{})
 	var instructionAvailable atomic.Value
 	// The tool sets instruction after execution.
 	tool := fantasy.NewAgentTool(
@@ -1624,14 +1624,16 @@ func TestRun_PrepareMessagesInjectsSystemContextMidLoop(t *testing.T) {
 			return nil
 		},
 		PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message {
-			if instructionInjected.Load() {
+			select {
+			case <-instructionInjected:
 				return nil
+			default:
 			}
 			instr, ok := instructionAvailable.Load().(string)
 			if !ok || instr == "" {
 				return nil
 			}
-			instructionInjected.Store(true)
+			close(instructionInjected)
 			// Insert a system message after existing system messages.
 			result := make([]fantasy.Message, 0, len(msgs)+1)
 			inserted := false
@@ -1748,7 +1750,7 @@ func TestRun_PrepareMessagesOnlyFiresOnce(t *testing.T) {
 		},
 	})
 	require.NoError(t, err)
-	require.Equal(t, 3, int(streamCalls))
+	require.Equal(t, 3, streamCalls)
 	// PrepareMessages is called before each of the 3 steps.
 	require.Equal(t, 3, int(prepareCalls.Load()))
 }
diff --git a/coderd/x/chatd/chattool/createworkspace_test.go b/coderd/x/chatd/chattool/createworkspace_test.go
@@ -1119,7 +1119,6 @@ func TestCreateWorkspace_OnChatUpdatedFiresAfterBuild(t *testing.T) {
 		GetChatWorkspaceTTL(gomock.Any()).
 		Return("0s", nil)
 
-
 	// UpdateChatWorkspaceBinding — triggers first OnChatUpdated.
 	db.EXPECT().
 		UpdateChatWorkspaceBinding(gomock.Any(), gomock.Any()).
@@ -1162,7 +1161,6 @@ func TestCreateWorkspace_OnChatUpdatedFiresAfterBuild(t *testing.T) {
 	var mu sync.Mutex
 	var callbackChats []database.Chat
 
-
 	createFn := func(_ context.Context, _ uuid.UUID, req codersdk.CreateWorkspaceRequest) (codersdk.Workspace, error) {
 		return codersdk.Workspace{
 			ID:        workspaceID,
diff --git a/coderd/x/chatd/instruction_test.go b/coderd/x/chatd/instruction_test.go
@@ -266,4 +266,24 @@ func TestInstructionFromContextFiles(t *testing.T) {
 		got := instructionFromContextFiles(msgs)
 		require.Empty(t, got)
 	})
+
+	t.Run("ReconstructsFromContextFileParts", func(t *testing.T) {
+		t.Parallel()
+		msgs := []database.ChatMessage{
+			makeMsg([]codersdk.ChatMessagePart{
+				{
+					Type:                 codersdk.ChatMessagePartTypeContextFile,
+					ContextFileOS:        "linux",
+					ContextFileDirectory: "/home/coder/project",
+					ContextFileContent:   "project rules",
+					ContextFilePath:      "/home/coder/project/AGENTS.md",
+				},
+			}),
+		}
+		got := instructionFromContextFiles(msgs)
+		require.Contains(t, got, "Operating System: linux")
+		require.Contains(t, got, "Working Directory: /home/coder/project")
+		require.Contains(t, got, "Source: /home/coder/project/AGENTS.md")
+		require.Contains(t, got, "project rules")
+	})
 }
