fix: re-fetch context files and skills from workspace on each turn #24360
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -4730,8 +4730,10 @@ func (p *Server) runChat( | |
| } | ||
| } | ||
|
|
||
| var instructionInjected bool | ||
| if instruction != "" { | ||
| prompt = chatprompt.InsertSystem(prompt, instruction) | ||
| instructionInjected = true | ||
| } | ||
| prompt = renderPlanPathPrompt(prompt, resolvePlanPathBlock(ctx)) | ||
| if skillIndex := chattool.FormatSkillIndex(skills); skillIndex != "" { | ||
|
|
@@ -5077,6 +5079,33 @@ func (p *Server) runChat( | |
| // start streaming build logs before the tool | ||
| // completes. | ||
| p.publishChatPubsubEvent(updatedChat, codersdk.ChatWatchEventKindStatusChange, nil) | ||
|
|
||
| // When a workspace is first attached mid-turn | ||
| // (e.g. via create_workspace), fetch and persist | ||
| // instruction files immediately so the LLM has | ||
| // AGENTS.md context for the remainder of this | ||
| // turn. The persisted marker prevents redundant | ||
| // fetches on subsequent turns. | ||
| if instruction == "" && updatedChat.WorkspaceID.Valid { | ||
| newInstruction, discoveredSkills, persistErr := p.persistInstructionFiles( | ||
| ctx, | ||
| updatedChat, | ||
| modelConfig.ID, | ||
| workspaceCtx.getWorkspaceAgent, | ||
| workspaceCtx.getWorkspaceConn, | ||
| ) | ||
| if persistErr != nil { | ||
| p.logger.Warn(ctx, "failed to persist instruction files on workspace attach", | ||
| slog.F("chat_id", updatedChat.ID), | ||
| slog.Error(persistErr), | ||
| ) | ||
| } else { | ||
| instruction = newInstruction | ||
| if len(discoveredSkills) > 0 { | ||
| skills = discoveredSkills | ||
| } | ||
| } | ||
| } | ||
| } | ||
| tools = append(tools, | ||
| chattool.ListTemplates(chat.OrganizationID, p.db, chattool.ListTemplatesOptions{ | ||
|
|
@@ -5311,11 +5340,27 @@ func (p *Server) runChat( | |
| if chat.ParentChatID.Valid { | ||
| reloadedPrompt = chatprompt.InsertSystem(reloadedPrompt, defaultSubagentInstruction) | ||
| } | ||
| // Re-derive instruction and skills from the reloaded | ||
| // messages so that any context added during the | ||
| // chatloop (e.g. via persistInstructionFiles when | ||
| // the agent changes) is picked up after compaction. | ||
| // The captured instruction takes priority; fall | ||
| // back to persisted DB content otherwise. | ||
| reloadedInstruction := instruction | ||
| if reloadedInstruction == "" { | ||
|
Member
There was a problem hiding this comment. P2 Double injection of instruction after compaction when instruction is acquired mid-turn (Edge Case Analyst, Contract Auditor) When This only manifests when compaction triggers on the same step as a workspace-creating tool call (conversation near context limit), so it's uncommon but not unrealistic for long-running chats. Fix: set
Member
Author
There was a problem hiding this comment. Fixed in a132d45 — added
Member
Author
There was a problem hiding this comment. Good catch — fixed in a132d45.
|
||
| reloadedInstruction = instructionFromContextFiles(reloadedMsgs) | ||
| } | ||
| if reloadedInstruction != "" { | ||
| reloadedPrompt = chatprompt.InsertSystem(reloadedPrompt, reloadedInstruction) | ||
| instructionInjected = true | ||
| } | ||
| reloadedPrompt = renderPlanPathPrompt(reloadedPrompt, resolvePlanPathBlock(reloadCtx)) | ||
| reloadedSkills := skillsFromParts(reloadedMsgs) | ||
| if len(reloadedSkills) == 0 { | ||
| reloadedSkills = skills | ||
| } | ||
|
|
||
| if skillIndex := chattool.FormatSkillIndex(reloadedSkills); skillIndex != "" { | ||
| reloadedPrompt = chatprompt.InsertSystem(reloadedPrompt, skillIndex) | ||
| } | ||
| reloadUserPrompt := p.resolveUserPrompt(reloadCtx, chat.OwnerID) | ||
|
|
@@ -5333,7 +5378,17 @@ func (p *Server) runChat( | |
| 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, | ||
|
|
@@ -5726,39 +5781,37 @@ func contextFileAgentID(messages []database.ChatMessage) (uuid.UUID, bool) { | |
| return lastID, found | ||
| } | ||
|
|
||
| // fetchWorkspaceContext retrieves fresh instruction files and | ||
| // skills from the workspace agent without persisting. It handles | ||
| // agent connection, context configuration fetching, content | ||
| // sanitization, and metadata stamping. Returns the workspace | ||
| // agent, the stamped parts, discovered skills, and whether the | ||
| // workspace connection succeeded. A nil agent means the chat has | ||
| // no valid workspace or the agent lookup failed; | ||
| // workspaceConnOK is false in that case. | ||
| func (p *Server) fetchWorkspaceContext( | ||
| ctx context.Context, | ||
|
Member
There was a problem hiding this comment. P2
The function signals failures via nil agent / false Carried forward from round 1 (Obs, no response) — upgrading to P2 given convergence across three reviewers.
Member
Author
There was a problem hiding this comment. Dropped the
|
||
| chat database.Chat, | ||
| getWorkspaceAgent func(context.Context) (database.WorkspaceAgent, error), | ||
| getWorkspaceConn func(context.Context) (workspacesdk.AgentConn, error), | ||
| ) (agent *database.WorkspaceAgent, agentParts []codersdk.ChatMessagePart, discoveredSkills []chattool.SkillMeta, workspaceConnOK bool) { | ||
| if !chat.WorkspaceID.Valid || getWorkspaceAgent == nil { | ||
| return nil, nil, nil, false | ||
| } | ||
|
|
||
| loadedAgent, agentErr := getWorkspaceAgent(ctx) | ||
| if agentErr != nil { | ||
| return nil, nil, nil, false | ||
| } | ||
|
|
||
| directory := loadedAgent.ExpandedDirectory | ||
| if directory == "" { | ||
| directory = loadedAgent.Directory | ||
| } | ||
|
|
||
| // Fetch context configuration from the agent. Parts | ||
| // arrive pre-populated with context-file and skill entries | ||
| // so we don't need additional round-trips. | ||
| if getWorkspaceConn != nil { | ||
| instructionCtx, cancel := context.WithTimeout(ctx, p.instructionLookupTimeout) | ||
| defer cancel() | ||
|
|
@@ -5789,21 +5842,15 @@ func (p *Server) persistInstructionFiles( | |
| // Stamp server-side fields and sanitize content. The | ||
| // agent cannot know its own UUID, OS metadata, or | ||
| // directory — those are added here at the trust boundary. | ||
| agentID := uuid.NullUUID{UUID: loadedAgent.ID, Valid: true} | ||
|
|
||
| for i := range agentParts { | ||
| agentParts[i].ContextFileAgentID = agentID | ||
| switch agentParts[i].Type { | ||
| case codersdk.ChatMessagePartTypeContextFile: | ||
| agentParts[i].ContextFileContent = SanitizePromptText(agentParts[i].ContextFileContent) | ||
| agentParts[i].ContextFileOS = loadedAgent.OperatingSystem | ||
| agentParts[i].ContextFileDirectory = directory | ||
| case codersdk.ChatMessagePartTypeSkill: | ||
| discoveredSkills = append(discoveredSkills, chattool.SkillMeta{ | ||
| Name: agentParts[i].SkillName, | ||
|
|
@@ -5814,6 +5861,49 @@ func (p *Server) persistInstructionFiles( | |
| } | ||
| } | ||
|
|
||
| return &loadedAgent, agentParts, discoveredSkills, workspaceConnOK | ||
| } | ||
|
|
||
| // persistInstructionFiles fetches AGENTS.md instruction files and | ||
| // skills from the workspace agent, persisting both as message | ||
| // parts. This is called once when a workspace is first attached | ||
| // to a chat (or when the agent changes). Returns the formatted | ||
| // instruction string and skill index for injection into the | ||
| // current turn's prompt. | ||
| func (p *Server) persistInstructionFiles( | ||
| ctx context.Context, | ||
| chat database.Chat, | ||
| modelConfigID uuid.UUID, | ||
| getWorkspaceAgent func(context.Context) (database.WorkspaceAgent, error), | ||
| getWorkspaceConn func(context.Context) (workspacesdk.AgentConn, error), | ||
| ) (instruction string, skills []chattool.SkillMeta, err error) { | ||
| agent, agentParts, discoveredSkills, workspaceConnOK := p.fetchWorkspaceContext( | ||
| ctx, chat, getWorkspaceAgent, getWorkspaceConn, | ||
| ) | ||
| // Defensive guard: fetchWorkspaceContext returns nil when the | ||
| // chat has no valid workspace or the agent lookup fails. It's | ||
| // cheaper to guard here than push the precondition up to all | ||
| // callers. | ||
| if agent == nil { | ||
| return "", nil, nil | ||
|
Member
There was a problem hiding this comment. we shouldn't be calling this if there is no workspace agent?
Member
Author
There was a problem hiding this comment. Added a comment explaining this is a defensive guard.
|
||
| } | ||
|
|
||
| agentID := uuid.NullUUID{UUID: agent.ID, Valid: true} | ||
| hasContent := false | ||
| hasContextFilePart := false | ||
| for _, part := range agentParts { | ||
| if part.Type == codersdk.ChatMessagePartTypeContextFile { | ||
| hasContextFilePart = true | ||
| if part.ContextFileContent != "" { | ||
| hasContent = true | ||
| } | ||
| } | ||
| } | ||
| directory := agent.ExpandedDirectory | ||
| if directory == "" { | ||
| directory = agent.Directory | ||
| } | ||
|
|
||
| if !hasContent { | ||
| if !workspaceConnOK { | ||
| return "", nil, nil | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
P2
onChatUpdatedfires beforewaitForBuild— mid-turn persist is a no-op for new workspaces (Edge Case Analyst P2, Contract Auditor P2)The approved plan addresses this: step 2d adds a second
OnChatUpdatedcall afterwaitForBuildin bothcreateworkspace.goandstartworkspace.go, and aPrepareMessageschatloop callback to inject the instruction into the running messages. Will implement in the next commit.