Skip to content

Commit ee60ee3

Browse files
hugodutkahugodutka
committed
fix(coderd/x/chatd): always commit a workspace context marker
1 parent a5a4c49 commit ee60ee3

3 files changed

Lines changed: 124 additions & 18 deletions

File tree

‎coderd/x/chatd/chatd_test.go‎

Lines changed: 88 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12693,8 +12693,8 @@ func TestActiveServer_WorkspaceContextAndDynamicToolInjection(t *testing.T) {
1269312693
require.Equal(t, database.ChatStatusWaiting, chatResult.Status)
1269412694

1269512695
parts := persistedChatParts(ctx, t, db, chat.ID)
12696-
require.Len(t, contextFilePartsForAgent(parts, dbAgent.ID), 1)
12697-
contextPart := contextFilePartsForAgent(parts, dbAgent.ID)[0]
12696+
require.Len(t, allContextFilePartsForAgent(parts, dbAgent.ID), 1)
12697+
contextPart := allContextFilePartsForAgent(parts, dbAgent.ID)[0]
1269812698
require.Equal(t, "/home/coder/project/AGENTS.md", contextPart.ContextFilePath)
1269912699
require.Equal(t, contextText, contextPart.ContextFileContent)
1270012700
require.Equal(t, "linux", contextPart.ContextFileOS)
@@ -12785,7 +12785,7 @@ func TestActiveServer_WorkspaceContextAndDynamicToolInjection(t *testing.T) {
1278512785
require.Equal(t, database.ChatStatusWaiting, secondResult.Status)
1278612786

1278712787
parts := persistedChatParts(ctx, t, db, chat.ID)
12788-
require.Len(t, contextFilePartsForAgent(parts, dbAgent.ID), 1)
12788+
require.Len(t, allContextFilePartsForAgent(parts, dbAgent.ID), 1)
1278912789
require.Equal(t, int32(1), contextConfigCalls.Load())
1279012790

1279112791
requestsMu.Lock()
@@ -12796,6 +12796,72 @@ func TestActiveServer_WorkspaceContextAndDynamicToolInjection(t *testing.T) {
1279612796
require.True(t, requestHasSystemSubstring(recorded[len(recorded)-1], contextText))
1279712797
})
1279812798

12799+
t.Run("commits marker when selected agent is unreachable", func(t *testing.T) {
12800+
t.Parallel()
12801+
12802+
db, ps := dbtestutil.NewDB(t)
12803+
ctx := testutil.Context(t, testutil.WaitLong)
12804+
12805+
var (
12806+
requestsMu sync.Mutex
12807+
requests []recordedOpenAIRequest
12808+
)
12809+
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
12810+
if !req.Stream {
12811+
return chattest.OpenAINonStreamingResponse("title")
12812+
}
12813+
12814+
requestsMu.Lock()
12815+
requests = append(requests, recordOpenAIRequest(req))
12816+
requestsMu.Unlock()
12817+
12818+
return chattest.OpenAIStreamingResponse(
12819+
chattest.OpenAITextChunks("done")...,
12820+
)
12821+
})
12822+
12823+
user, org, model := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL)
12824+
ws, dbAgent := seedWorkspaceWithAgent(t, db, user.ID)
12825+
require.NoError(t, db.SoftDeleteWorkspaceAgentsByWorkspaceID(ctx, ws.ID))
12826+
12827+
var agentDialCalls atomic.Int32
12828+
server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) {
12829+
cfg.AgentConn = func(context.Context, uuid.UUID) (workspacesdk.AgentConn, func(), error) {
12830+
agentDialCalls.Add(1)
12831+
return nil, nil, xerrors.New("unexpected workspace agent dial")
12832+
}
12833+
})
12834+
12835+
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
12836+
OrganizationID: org.ID,
12837+
OwnerID: user.ID,
12838+
APIKeyID: testAPIKeyID(t, db, user.ID),
12839+
Title: "workspace-context-agent-unreachable",
12840+
ModelConfigID: model.ID,
12841+
WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true},
12842+
AgentID: uuid.NullUUID{UUID: dbAgent.ID, Valid: true},
12843+
InitialUserContent: []codersdk.ChatMessagePart{
12844+
codersdk.ChatMessageText("Continue without workspace context."),
12845+
},
12846+
})
12847+
require.NoError(t, err)
12848+
12849+
chatResult := waitForTerminalChat(ctx, t, db, chat.ID)
12850+
require.Equal(t, database.ChatStatusWaiting, chatResult.Status)
12851+
12852+
parts := persistedChatParts(ctx, t, db, chat.ID)
12853+
markers := contextFileMarkersForAgent(parts, dbAgent.ID)
12854+
require.Len(t, markers, 1)
12855+
require.Empty(t, markers[0].ContextFileContent)
12856+
12857+
requestsMu.Lock()
12858+
recorded := append([]recordedOpenAIRequest(nil), requests...)
12859+
requestsMu.Unlock()
12860+
require.Len(t, recorded, 1, "expected model call after marker commit")
12861+
require.False(t, requestHasSystemSubstring(recorded[0], "Source: "))
12862+
require.Zero(t, agentDialCalls.Load())
12863+
})
12864+
1279912865
t.Run("repersists workspace context after agent changes", func(t *testing.T) {
1280012866
t.Parallel()
1280112867

@@ -12892,8 +12958,8 @@ func TestActiveServer_WorkspaceContextAndDynamicToolInjection(t *testing.T) {
1289212958
require.Equal(t, database.ChatStatusWaiting, secondResult.Status)
1289312959

1289412960
parts := persistedChatParts(ctx, t, db, chat.ID)
12895-
require.Len(t, contextFilePartsForAgent(parts, firstAgent.ID), 1)
12896-
require.Len(t, contextFilePartsForAgent(parts, secondAgent.ID), 1)
12961+
require.Len(t, allContextFilePartsForAgent(parts, firstAgent.ID), 1)
12962+
require.Len(t, allContextFilePartsForAgent(parts, secondAgent.ID), 1)
1289712963

1289812964
requestsMu.Lock()
1289912965
recorded := append([]recordedOpenAIRequest(nil), requests...)
@@ -12979,7 +13045,7 @@ func persistedChatMessages(
1297913045
return messages
1298013046
}
1298113047

12982-
func contextFilePartsForAgent(
13048+
func allContextFilePartsForAgent(
1298313049
parts []codersdk.ChatMessagePart,
1298413050
agentID uuid.UUID,
1298513051
) []codersdk.ChatMessagePart {
@@ -12996,6 +13062,22 @@ func contextFilePartsForAgent(
1299613062
return matched
1299713063
}
1299813064

13065+
func contextFileMarkersForAgent(
13066+
parts []codersdk.ChatMessagePart,
13067+
agentID uuid.UUID,
13068+
) []codersdk.ChatMessagePart {
13069+
var matched []codersdk.ChatMessagePart
13070+
for _, part := range parts {
13071+
if part.Type != codersdk.ChatMessagePartTypeContextFile ||
13072+
!part.ContextFileAgentID.Valid ||
13073+
part.ContextFileAgentID.UUID != agentID {
13074+
continue
13075+
}
13076+
matched = append(matched, part)
13077+
}
13078+
return matched
13079+
}
13080+
1299913081
func requireChatToolPart(
1300013082
t *testing.T,
1300113083
messages []database.ChatMessage,

‎coderd/x/chatd/generation.go‎

Lines changed: 34 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -800,9 +800,9 @@ func compactionModel(opts chatloop.GenerateCompactionOptions) string {
800800
// workspace context messages (e.g. AGENTS.md, workspace skills) into
801801
// chat history. It records a generation attempt, calls the injected
802802
// workspace context builder without holding the DB lock, then commits
803-
// the returned messages fenced to the attempt. If the builder returns
804-
// no messages, the action exits as expected and the next worker task
805-
// re-reads the chat.
803+
// the returned messages fenced to the attempt. If context cannot be
804+
// fetched, it commits a marker for the selected agent so the generation
805+
// loop can continue.
806806
func (s *taskStarter) persistWorkspaceContext(
807807
ctx context.Context,
808808
machine *chatstate.ChatMachine,
@@ -831,20 +831,45 @@ func (s *taskStarter) persistWorkspaceContext(
831831
ActiveAPIKeyID: modelOpts.ActiveAPIKeyID,
832832
})
833833
if err != nil {
834-
if errors.Is(err, errWorkspaceContextUnavailable) {
835-
// Builder reported nothing durable to commit (workspace or
836-
// agent missing, unreachable, etc.). Exit the action without
837-
// committing so the next worker task can re-read the chat.
838-
return errTaskExpectedExit
834+
s.opts.Logger.Warn(ctx, "failed to build workspace context, committing marker",
835+
slog.F("chat_id", input.ChatID),
836+
slog.F("worker_id", input.WorkerID),
837+
slogError(err),
838+
)
839+
marker, err := workspaceContextMarkerMessage(locked, modelOpts.ActiveAPIKeyID)
840+
if err != nil {
841+
return xerrors.Errorf("build workspace context marker: %w", err)
839842
}
840-
return err
843+
result.Messages = []chatstate.Message{marker}
841844
}
842845
return s.commitGenerationStep(ctx, machine, input, attempt, generationActionPersistWorkspaceContext, stepMessagesForCommit{
843846
Messages: result.Messages,
844847
VisibleIndexes: visibleMessageIndexes(result.Messages),
845848
})
846849
}
847850

851+
// workspaceContextMarkerMessage builds an empty context-file sentinel
852+
// for the chat's selected agent. Committing this marker lets the
853+
// generation loop proceed when the agent is unreachable.
854+
func workspaceContextMarkerMessage(chat database.Chat, activeAPIKeyID string) (chatstate.Message, error) {
855+
content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{{
856+
Type: codersdk.ChatMessagePartTypeContextFile,
857+
ContextFileAgentID: chat.AgentID,
858+
}})
859+
if err != nil {
860+
return chatstate.Message{}, xerrors.Errorf("marshal workspace context marker: %w", err)
861+
}
862+
modelConfigID := chat.LastModelConfigID
863+
return chatstate.Message{
864+
Role: database.ChatMessageRoleUser,
865+
Content: content,
866+
Visibility: database.ChatMessageVisibilityBoth,
867+
ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: modelConfigID != uuid.Nil},
868+
ContentVersion: chatprompt.CurrentContentVersion,
869+
APIKeyID: sql.NullString{String: activeAPIKeyID, Valid: activeAPIKeyID != ""},
870+
}, nil
871+
}
872+
848873
func (s *taskStarter) beginGenerationAttempt(
849874
ctx context.Context,
850875
machine *chatstate.ChatMachine,

‎coderd/x/chatd/workspace_context_builder.go‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import (
1919
// errWorkspaceContextUnavailable is returned by buildWorkspaceContext
2020
// when there is nothing safe to persist for the current committed
2121
// metadata, e.g. the chat has no bound workspace agent or the agent is
22-
// no longer resolvable. Callers treat it as an expected exit.
22+
// no longer resolvable.
2323
var errWorkspaceContextUnavailable = xerrors.New("workspace context unavailable")
2424

2525
// buildWorkspaceContext fetches workspace context for the chat's
@@ -54,8 +54,7 @@ func (server *Server) buildWorkspaceContext(
5454
defer wsCtx.close()
5555

5656
parts, expectedAgentID := server.fetchContextForBuild(ctx, chat, &wsCtx, logger)
57-
// If the workspace or agent is gone, fall back to no-op so the
58-
// generation action exits without committing stale context.
57+
// If the workspace or agent is gone, report unavailable.
5958
if expectedAgentID == uuid.Nil {
6059
return workspaceContextBuildResult{}, errWorkspaceContextUnavailable
6160
}

0 commit comments

Comments
 (0)