coder/coderd/x/chatd/workspace_context_builder.go

package chatd

import (
	"context"
	"sync"

	"github.com/google/uuid"

	"cdr.dev/slog/v3"
	"github.com/coder/coder/v2/coderd/database"
	"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
	"github.com/coder/coder/v2/coderd/x/chatd/chatstate"
	"github.com/coder/coder/v2/codersdk"
	"github.com/coder/coder/v2/codersdk/workspacesdk"
)

// buildWorkspaceContext fetches workspace context for the chat's
// bound workspace agent and returns durable chatstate.Message values
// for the generation action to commit. It returns an empty result
// (Messages == nil) when there is nothing safe to persist for the
// current committed metadata; that is treated as an expected exit by
// the generation action.
func (server *Server) buildWorkspaceContext(
	ctx context.Context,
	input workspaceContextBuildInput,
) (workspaceContextBuildResult, error) {
	chat := input.Chat
	if !chat.WorkspaceID.Valid || !chat.AgentID.Valid {
		return workspaceContextBuildResult{}, nil
	}
	logger := server.logger.With(
		slog.F("chat_id", chat.ID),
		slog.F("owner_id", chat.OwnerID),
	)

	// Build a per-call workspace context with the latest committed
	// chat snapshot so getWorkspaceAgent and getWorkspaceConn dial
	// the agent we actually want to fetch context from.
	currentChat := chat
	var chatStateMu sync.Mutex
	wsCtx := turnWorkspaceContext{
		server:           server,
		chatStateMu:      &chatStateMu,
		currentChat:      &currentChat,
		loadChatSnapshot: server.db.GetChatByID,
	}
	defer wsCtx.close()

	parts, discoveredSkillsPart, _, expectedAgentID := server.fetchContextForBuild(ctx, chat, &wsCtx, logger)
	_ = discoveredSkillsPart
	// If the workspace or agent is gone, fall back to no-op so the
	// generation action exits without committing stale context.
	if expectedAgentID == uuid.Nil {
		return workspaceContextBuildResult{}, nil
	}

	hasContent := false
	hasContextFilePart := false
	for _, part := range parts {
		if part.Type == codersdk.ChatMessagePartTypeContextFile {
			hasContextFilePart = true
			if part.ContextFileContent != "" {
				hasContent = true
			}
		}
	}

	agentID := uuid.NullUUID{UUID: expectedAgentID, Valid: true}

	// If we have no content but the agent is known, commit a blank
	// context-file marker (sentinel) so subsequent turns skip the
	// workspace-agent dial and the decision helper observes the
	// attempt in committed history. This applies whether the
	// workspace connection succeeded but returned no AGENTS.md, or
	// the agent's context config fetch failed: in both cases we
	// have a known agent and committing a sentinel breaks the
	// otherwise-infinite decision loop.
	if !hasContent {
		if !hasContextFilePart {
			parts = append([]codersdk.ChatMessagePart{{
				Type:               codersdk.ChatMessagePartTypeContextFile,
				ContextFileAgentID: agentID,
			}}, parts...)
		}
	}

	content, err := chatprompt.MarshalParts(parts)
	if err != nil {
		return workspaceContextBuildResult{}, nil
	}

	modelConfigID := chat.LastModelConfigID
	msg := chatstate.Message{
		Role:           database.ChatMessageRoleUser,
		Content:        content,
		Visibility:     database.ChatMessageVisibilityBoth,
		ModelConfigID:  uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil},
		ContentVersion: chatprompt.CurrentContentVersion,
	}

	// Update the cache column so subsequent turns can read the last
	// injected context without scanning messages. This is a
	// best-effort write that does not mutate chat history; the
	// generation action separately commits the durable message
	// below.
	stripped := make([]codersdk.ChatMessagePart, len(parts))
	copy(stripped, parts)
	for i := range stripped {
		stripped[i].StripInternal()
	}
	server.updateLastInjectedContext(ctx, chat.ID, stripped)

	return workspaceContextBuildResult{Messages: []chatstate.Message{msg}}, nil
}

// fetchContextForBuild fetches workspace context parts from the
// agent, returning the parts to persist, any discovered skill
// metadata (carried in parts), and whether the workspace connection
// succeeded. expectedAgentID is the agent ID the fetch was bound to,
// or uuid.Nil if the agent could not be resolved.
func (server *Server) fetchContextForBuild(
	ctx context.Context,
	chat database.Chat,
	wsCtx *turnWorkspaceContext,
	logger slog.Logger,
) (parts []codersdk.ChatMessagePart, discoveredSkills []codersdk.ChatMessagePart, workspaceConnOK bool, expectedAgentID uuid.UUID) {
	agent, agentParts, _, connOK := server.fetchWorkspaceContext(
		ctx, chat, wsCtx.getWorkspaceAgent,
		func(instructionCtx context.Context) (workspacesdk.AgentConn, error) {
			if _, _, err := wsCtx.workspaceAgentIDForConn(instructionCtx); err != nil {
				return nil, err
			}
			return wsCtx.getWorkspaceConn(instructionCtx)
		},
	)
	if agent == nil {
		// fetchWorkspaceContext returns nil for the agent when the
		// chat has no valid workspace or the agent lookup fails.
		logger.Debug(ctx, "workspace context build: workspace agent not resolvable")
		return nil, nil, false, uuid.Nil
	}
	return agentParts, agentParts, connOK, agent.ID
}