mirror of
https://github.com/coder/coder.git
synced 2026-09-21 20:51:01 +08:00
feat(coderd): wire advisor runtime to admin config (#24622)
## Summary Wire the advisor runtime into `chatd`: read the admin config on every `runChat`, gate tool registration and system-prompt guidance on a **single eligibility boolean**, register the `advisor` built-in tool, and apply the exclusive-tool policy from PR 1. ## Motivation This is the integration seam where PRs 1–3 come together into an actual user-visible feature. Gating is deliberately root-chat-only for the initial rollout; child/sub-agent chats still do not see the tool or the guidance block. ## Changes ### `coderd/x/chatd/chatd.go` - `loadAdvisorConfig(ctx, logger)` reads the admin config (from PR 3) on each run. If `ModelConfigID` is set, it resolves the override model via `configCache.ModelConfigByID`; otherwise it falls back to the outer chat's model and provider options. Reasoning effort is plumbed into provider options via `applyAdvisorReasoningEffort`. - One computed `advisorEligible` boolean drives **both** tool registration (after skill tools, before MCP tools) and guidance injection via `chatprompt.InsertSystem(prompt, chatadvisor.ParentGuidanceBlock)`. - `setAdvisorPromptSnapshot` closures capture the outer prompt state at the right points in the lifecycle (`renderPlanPathPrompt`, `ReloadMessages`, `PrepareMessages`) so the advisor handoff uses the same context the outer model saw. - `ExclusiveToolNames["advisor"] = true` is passed to `chatloop.Run()` so mixed batches are rejected cleanly (PR 1 machinery). - `builtinToolNames["advisor"] = true` so metrics keep advisor distinct from the generic `mcp` label. ### Child-chat guard - Child/sub-agent chats deliberately do not see the advisor tool or guidance block, to avoid recursion/cost blowups until the pattern is proven. This is covered by `TestAdvisorGating_ChildChat` (currently skipped pending a rewrite against the new `plan`/`explore` subagent infrastructure; core gating logic is still exercised by `TestAdvisorGating_Disabled` and `TestAdvisorGating_RootChat`). ## Stack context This is **PR 4 of 6** in the advisor feature stack. It depends on PRs 1–3. ## Scope / non-goals - No frontend changes. The feature is invocable via the backend but renders generically until PR 5. - No separate provider runner; the nested advisor call reuses the existing model/provider path. - No DB migration. ## Validation - `go test ./coderd/x/chatd/... -run TestAdvisor` - `go build ./...` - `make lint` --- <details> <summary>📋 Implementation Plan (shared across the advisor stack)</summary> # Plan: Add a Mux-style advisor tool to coder agents/chatd ## Outcome Add a first-class `advisor` tool to agent chats in `coderd/x/chatd` that feels native to Coder: - it is a built-in server-side tool, not an MCP/dynamic-tool workaround; - it performs a nested **tool-less** model call for strategic advice; - it is exposed only when eligible, and the prompt mentions it only when it is actually available; - it is treated as a **planning-only** tool so it does not run alongside action tools in the same batch; - it tracks usage/cost separately enough for operators to reason about it; - it has a minimally polished UI in the Agents page; - and it ships with explicit dogfooding evidence, including screenshots and repro videos. ## Design decisions to lock before coding 1. **Primary architecture:** native built-in tool in `chattool/`, backed by a small `chatadvisor` package. 2. **Nested model execution:** reuse chatd's existing model/provider stack for a one-step, tool-less advisor call rather than inventing a new provider pathway. 3. **Execution policy:** treat `advisor` as an exclusive/planning-only tool; mixed batches must return structured policy errors and force the model to retry cleanly. 4. **Availability:** initial rollout is for root agent chats only; disable for child/sub-agent chats until recursion/cost policy is proven. 5. **Prompt sync:** use one eligibility boolean to drive both tool registration and advisor guidance injection. 6. **Persistence/cost split:** MVP should keep advisor usage visible in result metadata and server metrics; only add DB schema if product/billing explicitly needs queryable advisor-specific cost. 7. **UI scope:** generic tool rendering is an acceptable temporary milestone during backend bring-up, but the release candidate should include a dedicated lightweight advisor renderer. ## Delivery model The work should be executed as coordinated workstreams with one integration owner and parallel contributors for low-conflict areas. The integration owner should own `coderd/x/chatd/chatd.go` because prompt assembly, tool registration, and model resolution all converge there. ## Detailed workstreams ### Repo evidence used for this plan <details> <summary>Mux reference and current chatd seams</summary> **Mux reference implementation** - `src/node/services/tools/advisor.ts` — native advisor tool implementation. - `src/common/constants/advisor.ts` — advisor prompt/constants and truncation policy. - `src/common/utils/tools/tools.ts` — conditional tool registration. - `src/node/services/streamContextBuilder.ts` — injects advisor guidance only when the tool is available. **Current chatd seams** - `coderd/x/chatd/chatd.go` - `processChat()` — tool assembly, prompt assembly, and chatloop invocation. - `resolveChatModel()` — current model/provider/key resolution seam. - `type Config struct` — server-level chatd configuration surface. - `coderd/x/chatd/chatloop/chatloop.go` - `Run()` — main streaming/model loop. - `executeTools()` — built-in tool execution/batching seam. - `coderd/x/chatd/chattool/` — built-in tool implementations. - `site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx` — tool renderer dispatch. - `site/src/pages/AgentsPage/components/ChatConversation/messageParsing.ts` and `ConversationTimeline.tsx` — tool/result merge and rendering flow. </details> ### Workstream map and ownership | Workstream | Primary owner | Main files | Can run in parallel? | Done when | |---|---|---|---|---| | 0. Integration + gating | Integration lead | `coderd/x/chatd/chatd.go` | No; central merge lane | Tool registration, prompt sync, and model selection are wired together | | 1. Advisor runtime + tool | Backend agent | new `coderd/x/chatd/chatadvisor/`, new `coderd/x/chatd/chattool/advisor.go` | Yes | Tool can perform a tool-less advisor call in memory and return structured results | | 2. Planning-only execution policy | Chatloop agent | `coderd/x/chatd/chatloop/chatloop.go`, related tests | Yes | Mixed `advisor` + action-tool batches are rejected cleanly and deterministically | | 3. Metrics/usage/config | Backend/telemetry agent | `chatd.go`, `chatloop/metrics.go`, optional config plumbing | Partially; coordinate with integration lead | Advisor usage is separately visible in metadata/metrics and limits are enforced | | 4. Frontend rendering | Frontend agent | `site/.../tools/Tool.tsx`, new `AdvisorTool.tsx`, stories | Yes after result schema stabilizes | Advisor renders as a readable card and story tests pass | | 5. Dogfood + QA evidence | QA agent | dev server, Storybook, dogfood output | After backend + UI are usable | Repro videos, screenshots, and a concise QA report exist | ### Parallelization rules - **Do not split `coderd/x/chatd/chatd.go` across multiple execution agents without an integration lead.** That file owns prompt building, tool registration, model resolution, and cost persistence. - Workstreams 1 and 2 can be developed in parallel and then stacked onto the integration branch. - Workstream 4 should begin once the backend result schema is agreed on, even if the backend is still behind a feature flag. - Any agent that needs to re-check Mux behavior should clone `coder/mux` into a temporary directory (for example, `$(mktemp -d)/mux`) and inspect it read-only; do not vendor or copy code from Mux directly. ## Phase 0 — Preflight and guardrails ### Goals - Align the team on the smallest shippable architecture. - Prevent scope creep into MCP/dynamic-tool/sub-agent variants. - Decide upfront what is MVP vs. follow-up. ### Tasks 1. **Confirm the MVP boundary.** - Ship a built-in advisor tool first. - Do **not** make MCP, dynamic tools, or sub-agents the primary implementation. - Do **not** add transient streaming phases in the first backend PR unless they fall out almost for free. 2. **Confirm local workflow hygiene before coding.** - Ensure the repo is using the project git hooks from `scripts/githooks`. - Do not bypass hooks with `--no-verify`. - Use `./scripts/develop.sh` for the full dev server rather than manual build/run commands. 3. **Lock the model-selection policy.** - **Recommended MVP:** advisor uses the same resolved provider/model/cost config as the current chat, with advisor-specific max-output and usage caps. - **Follow-up only if required:** add a separate `AdvisorModelConfigID`-style override that resolves through the existing `configCache`/model-config path. Do not invent a new free-form `provider:model` parser if chatd already stores provider/model separately. 4. **Lock the persistence policy.** - **Recommended MVP:** no DB migration. Persist advisor-visible metadata in the tool result and record separate metrics in memory/Prometheus. - **Only if product/billing explicitly asks for queryable advisor cost:** add a later DB migration or usage table, following the normal `queries/*.sql` + `make gen` workflow. 5. **Create an execution ADR note in the work item or tracking doc.** - Capture: built-in tool, tool-less nested call, root-chat-only rollout, exclusive execution policy, MVP no-DB-migration default. ### Quality gate - Everyone on the team can state the same answers to these questions: - Is advisor a built-in tool? **Yes.** - Can advisor run with action tools in the same batch? **No.** - Does advisor get tools of its own? **No.** - Is a DB migration required for MVP? **No, unless billing insists.** ## Phase 1 — Build the advisor runtime and tool wrapper ### Goals Create the core advisor implementation in a way that is easy to test and keeps `chattool/` thin. ### Files to add - `coderd/x/chatd/chatadvisor/types.go` - `coderd/x/chatd/chatadvisor/guidance.go` - `coderd/x/chatd/chatadvisor/handoff.go` - `coderd/x/chatd/chatadvisor/runtime.go` - `coderd/x/chatd/chatadvisor/runner.go` - `coderd/x/chatd/chattool/advisor.go` ### Responsibilities by file 1. **`types.go`** - Define the input/result schema used by the tool and UI. - Keep the result shape close to Mux so the UI and model both have predictable cases. - Recommended result variants: - `advice` - `limit_reached` - `error` Recommended shape: ```go type AdvisorArgs struct { Question string `json:"question"` } type AdvisorResult struct { Type string `json:"type"` Advice string `json:"advice,omitempty"` Error string `json:"error,omitempty"` AdvisorModel string `json:"advisor_model,omitempty"` RemainingUses int `json:"remaining_uses,omitempty"` Usage *AdvisorUsageResult `json:"usage,omitempty"` } ``` 2. **`guidance.go`** - Hold two strings: - the nested advisor system prompt; - the parent-agent guidance block to inject into the outer system prompt. - The nested advisor prompt must say, in plain language: - you are advising the parent agent; - you do not address the end user directly; - you do not claim actions happened; - you return concise strategic guidance and tradeoffs. 3. **`runtime.go`** - Define the per-run runtime state. - Recommended fields: - resolved model + model config; - provider keys/options reused from the outer chat; - `MaxUsesPerRun`; - `MaxOutputTokens`; - atomic/current call counter; - callback(s) to obtain the current prompt snapshot and current-step snapshot; - optional metrics/usage hook. - Add fail-fast validation for impossible config: nil model, non-positive limits, empty prompt builders, etc. 4. **`handoff.go`** - Build the advisor handoff message from: - the explicit question; - the exact prompt/messages the parent model just used; - the current step's text/reasoning snapshot, if available; - the most recent relevant tool outputs, if they are already in the prompt snapshot. - **Important:** use the already-prepared outer prompt tail, not a fresh DB reload. That keeps the advisor aligned with compaction and the exact context the outer model saw. - Apply hard truncation budgets with recent-context bias. 5. **`runner.go`** - Execute the nested advisor call. - **Recommended implementation:** call `chatloop.Run()` in an in-memory, one-step mode: - `Tools: nil` - `ProviderTools: nil` - `MaxSteps: 1` - `PersistStep`: capture the assistant output in memory instead of writing DB rows - Reuse the existing provider/model/cost path instead of building a second provider runner. - Assert that no tool definitions are passed to the nested call. 6. **`chattool/advisor.go`** - Keep this file thin and consistent with other built-ins. - Responsibilities: - decode `AdvisorArgs`; - validate `Question` is non-empty and bounded; - call the `chatadvisor` runner; - return a structured tool response. ### Defensive programming requirements - Assert `Question` is non-empty after trimming. - Assert runtime limits are positive. - Assert the nested advisor call runs with zero tools/provider tools. - Assert `AdvisorResult.Type` is one of the known variants before returning. - Assert remaining uses never goes negative. ### Acceptance criteria - A unit test can call the advisor tool with a fake model and receive a stable `advice` result. - The nested advisor call is impossible to run with tools accidentally attached. - The core logic lives in `chatadvisor/`, not embedded inside `chatd.go`. ## Phase 2 — Wire advisor into chatd and keep prompt/tool availability in sync ### Goals Register the tool in the right place, expose it only when eligible, and inject system guidance only when the tool is present. ### Files to modify - `coderd/x/chatd/chatd.go` - optionally a small helper file if `chatd.go` becomes too crowded ### Tasks 1. **Compute one eligibility boolean in `processChat()`.** Recommended inputs: - server-level advisor enabled flag; - root chat only (`chat.ParentChatID == uuid.Nil` or equivalent existing root/child check); - a usable resolved model/provider exists; - optional experiment/workspace/org gate if product wants staged rollout. 2. **Create the runtime once per outer chat run.** - Use the model/config/keys resolved by `resolveChatModel()`. - Reuse provider options from the current chat's `ChatModelCallConfig`. - Set `MaxUsesPerRun` and `MaxOutputTokens` from advisor config defaults. 3. **Register the tool in the built-in tool block.** - Insert after the skill tools and before MCP tools in `processChat()`. - Record `builtinToolNames["advisor"] = true` so metrics stay bounded. 4. **Inject advisor guidance into the outer system prompt using the same boolean.** - Use `chatprompt.InsertSystem()` in the same prompt assembly path that already injects user/system instructions. - Place the block near the existing instruction insertion, before plan-path/skill context blocks. - Wrap the guidance in an explicit tag like `<advisor-guidance>` so it is easy to spot in tests and future refactors. 5. **Keep advisor out of child chats for the first release.** - That avoids recursion/cost blowups with `spawn_agent` / `wait_agent` flows. - Document this explicitly in the rollout notes and tests. ### Acceptance criteria - If advisor is disabled, neither the tool nor the prompt guidance appears. - If advisor is enabled, both the tool and the prompt guidance appear. - Root chats can use advisor; child chats cannot. - Built-in tool names include `advisor` so metrics do not collapse it into the generic `mcp` label. ## Phase 3 — Enforce planning-only execution policy in `chatloop` ### Goals Prevent the model from calling `advisor` and action tools in the same execution batch. ### Files to modify - `coderd/x/chatd/chatloop/chatloop.go` - related chatloop tests ### Recommended implementation Keep the MVP small; do **not** build a general policy engine yet. 1. Add a minimal field to `chatloop.RunOptions`, for example: ```go ExclusiveToolName *string ``` 2. In `Run()` / `executeTools()`, detect the case where the exclusive tool appears in the same local-tool batch as any other locally executed tool. 3. When that happens, synthesize structured tool-result errors for the affected calls instead of executing anything in the batch. - `advisor` should receive a clear error like: _advisor must be called by itself before action tools_. - The sibling action tools should receive a paired policy error like: _this tool was skipped because advisor must run alone_. 4. Let the outer model see those tool errors and retry cleanly. - This is simpler and safer than partial execution or hidden deferral. - It preserves deterministic transcript history for debugging. 5. Pass the just-finished step snapshot into the tool execution context. - The advisor runtime should be able to see the current step's text/reasoning content, because that is often the best hint about what the outer model is trying to decide. ### Why this is the right fit - It matches the intended semantics: advisor is consulted **before** taking action. - It avoids subtle race conditions caused by concurrent built-in tool execution. - It keeps the behavior easy to test with fake models. ### Acceptance criteria - A model-emitted batch containing only `advisor` succeeds. - A model-emitted batch containing `advisor` plus any other locally executed tool returns deterministic policy errors and executes nothing. - Non-advisor tool execution stays unchanged for normal chats. ## Phase 4 — Usage limits, metrics, and configuration ### Goals Make advisor safe to operate without over-designing billing/storage in the first release. ### Files to modify - `coderd/x/chatd/chatd.go` - `coderd/x/chatd/chatloop/metrics.go` as needed - `coderd/x/chatd/chatd.go` `Config` struct and constructor path - optional follow-up config/db files only if a separate advisor model or persistent billing is required ### Tasks 1. **Add explicit server config knobs for MVP.** Recommended fields on `chatd.Config` or a nested advisor config struct: - `AdvisorEnabled bool` - `AdvisorMaxUsesPerRun int` - `AdvisorMaxOutputTokens int64` 2. **Track usage per outer run.** - Reset the counter for each `processChat()` invocation. - Return `remaining_uses` in the tool result. - Return `limit_reached` when the cap is exhausted. 3. **Expose advisor usage metadata in the tool result.** - Include model name and token/cost summary if available. - Use the same `callConfig.Cost` calculation path as the outer chat for MVP if advisor reuses the same model. 4. **Record server-side metrics.** - Count advisor invocations, failures, and latency. - Ensure they show up under the built-in tool label `advisor`. 5. **Optional decision gate: separate advisor model.** - If product insists on a stronger/different advisor model, add a follow-up config hook that resolves another existing chat model config through the same `configCache` path. - Keep that out of the first landing PR unless it is required for acceptance. 6. **Optional decision gate: queryable advisor cost.** - If this becomes required, spin a follow-up DB task: - update `coderd/database/queries/*.sql`; - add migration files; - run `make gen`; - update audit mappings if a new auditable type/field is introduced. ### Acceptance criteria - Advisor calls are capped per outer run. - Limit exhaustion is user-visible in the tool result. - Metrics distinguish advisor calls from other built-in tools. - MVP does not require a schema migration unless explicitly approved. ## Phase 5 — Frontend rendering and Storybook coverage ### Goals Make advisor feel intentional in the Agents UI without blocking the backend on fancy streaming UI. ### Files to modify - `site/src/pages/AgentsPage/components/ChatElements/tools/Tool.tsx` - new `site/src/pages/AgentsPage/components/ChatElements/tools/AdvisorTool.tsx` - Storybook story file(s) in the same tools directory ### Delivery strategy 1. **Intermediate milestone during backend bring-up:** rely on the existing generic tool renderer if needed. - This is acceptable only as a short-lived integration checkpoint. 2. **Release milestone:** add a dedicated lightweight `AdvisorTool` renderer. - Reuse existing primitives: - `ToolCollapsible` - `ToolIcon` - `Response` for markdown/prose rendering - `ScrollArea` if the advice can be long - Keep styling light and consistent with the Agents page. - Do not add unnecessary React memoization in `site/src/pages/AgentsPage/`; that area is already React-Compiler aware. 3. **Render the structured result states cleanly.** - `advice` — readable prose/markdown with optional metadata footer. - `limit_reached` — warning-style message. - `error` — error state with visible fallback text. - `running` — existing tool loading state/spinner is enough for MVP. 4. **Add Storybook coverage instead of ad-hoc component tests.** Recommended stories: - successful advice; - running/loading; - limit reached; - error. 5. **Keep the UI contract narrow.** - Prefer one text field like `advice` plus small metadata rather than a deeply nested schema. - That keeps the UI resilient to prompt iteration. ### Acceptance criteria - The advisor tool card renders readable content rather than raw quoted JSON in the final release branch. - Running, limit, and error states are visibly distinct. - Storybook stories and play assertions cover the new states. - Existing tool rendering flows remain unchanged. ## Phase 6 — Automated tests and validation gates ### Backend tests to add 1. **Advisor runtime/tool tests** - question validation; - tool-less nested execution assertion; - success result shaping; - limit-reached result shaping; - error result shaping. 2. **Prompt/gating tests in chatd** - advisor disabled ⇒ no tool, no guidance; - advisor enabled/root chat ⇒ tool + guidance; - child chat ⇒ advisor absent. 3. **Chatloop policy tests** - advisor alone runs; - advisor + action tool mixed batch returns deterministic policy errors; - non-advisor tools still execute normally. 4. **Usage/metrics tests** - per-run cap resets correctly; - builtin tool labeling includes `advisor`; - returned metadata includes model/usage summary when available. ### Frontend tests to add - Storybook `play()` assertions for the advisor renderer states. - Verify expand/collapse behavior and visible fallback text. - Verify the message timeline still renders adjacent tools correctly. ### Recommended command sequence Run these as the implementation matures, not only at the end: 1. Backend-focused gate after phases 1–4: - `make test RUN=TestAdvisor` - `make test RUN=TestChatloopAdvisor` - `make lint` 2. Frontend-focused gate after phase 5: - `pnpm test:storybook src/pages/AgentsPage/components/ChatElements/tools/AdvisorTool.stories.tsx` - `pnpm lint` - `pnpm format` 3. Final repo gate before handoff: - `make pre-commit` - run any additional targeted `make test RUN=...` selections covering touched chatd paths > Use the exact new test names the implementing agents create; the names above are recommended anchors, not existing tests. ## Dogfooding plan ### Principle Dogfood the change as a real agent feature, not just a unit-tested backend. Per the dogfood and `agent-browser` skills, the reviewer should get **watchable repro videos** plus screenshots that make the behavior obvious without reading logs. ### Required setup 1. Start the full dev environment with: - `./scripts/develop.sh` 2. If the frontend renderer changes, also start Storybook from `site/` with: - `pnpm storybook --no-open` 3. Use `agent-browser` directly — **never `npx agent-browser`**. 4. Use named browser sessions and an output folder such as: - `./dogfood-output/advisor/` - with subfolders `screenshots/` and `videos/` ### Evidence protocol For every interactive scenario below: 1. Start video recording **before** the action. 2. Capture step-by-step screenshots at human pace. 3. Capture one annotated screenshot of the final state. 4. Stop the recording. 5. Note the exact pass/fail observation in the QA report. For static UI states (for example Storybook error/limit cards), an annotated screenshot is sufficient; video is optional but still encouraged by this project’s review preference. ### Dogfood scenarios #### Scenario A — Happy path in the real Agents UI **Goal:** prove that a root agent chat can invoke advisor and produce a readable recommendation before taking further action. Steps: 1. Open the Agents page with an advisor-enabled root chat. 2. Start a repro video. 3. Send a prompt that should reasonably trigger strategic planning, such as an architecture or multi-tradeoff question. 4. Capture screenshots of: - the prompt before send; - the running advisor state; - the completed advisor card and the assistant’s follow-up response. 5. Stop recording. Pass criteria: - advisor appears in the timeline; - the rendered result is readable; - the assistant can continue after consuming the advisor output. #### Scenario B — Advisor unavailable path **Goal:** prove the feature is truly gated. Suggested variants (at least one is required, both are better): - feature flag/config off; - child/sub-agent chat. Evidence: - annotated screenshot of the chat/tool state showing advisor is absent; - short video if toggling the gate live is part of the repro. Pass criteria: - no advisor tool is available; - no advisor-specific prompt behavior leaks through. #### Scenario C — UI states in Storybook **Goal:** prove the renderer handles non-happy states cleanly. Required story states: - success/advice; - running; - limit reached; - error. Evidence: - one screenshot per state; - at least one short video showing collapse/expand behavior. Pass criteria: - success renders readable advice; - limit/error have visible fallback text; - the component behaves like the other tool cards. #### Scenario D — Regression sweep of nearby tools **Goal:** ensure advisor does not break the surrounding chat timeline. Check at minimum: - another existing built-in tool still renders correctly near advisor; - sub-agent/tool cards still expand/collapse normally; - no obvious console errors appear in the Agents page during the advisor flow. Evidence: - screenshots of adjacent tool cards; - console/error capture if anything suspicious appears. ### `agent-browser` usage notes for the QA agent - Prefer `agent-browser batch` for 2+ sequential commands when no intermediate parsing is needed. - Use `snapshot -i` to discover interactive refs. - Re-snapshot after navigation or major DOM changes. - Avoid `wait --load networkidle` unless the page is known to go idle; prefer explicit element/text waits or short fixed waits. - Record videos at human pace and include pauses that a reviewer can follow. ## Rollout plan ### Initial rollout - Gate behind a server-side advisor-enabled flag. - Enable only for selected internal/root agent chats first. - Watch metrics for: - invocation count; - failure rate; - latency; - obvious retry loops. ### Expansion conditions Expand beyond the initial rollout only after the following are true: - mixed-batch policy behavior is stable; - cost impact is understood; - frontend UX is readable in production-like dogfood; - no recursion surprises have appeared with sub-agent flows. ### Explicit non-goals for the first release - advisor inside child/sub-agent chats; - provider-agnostic streaming phase UI; - MCP-based external advisor implementation; - mandatory DB-backed advisor cost reporting. ## Final acceptance checklist - [ ] `advisor` is a built-in chatd tool, not an MCP/dynamic-tool substitute. - [ ] The nested advisor call is tool-less and bounded to one in-memory step. - [ ] One eligibility boolean controls both tool registration and prompt guidance injection. - [ ] Root chats can use advisor; child chats cannot in the initial rollout. - [ ] Mixed advisor/action batches produce deterministic policy errors instead of partial execution. - [ ] Per-run usage caps and limit-reached behavior work. - [ ] Advisor usage is visible in metadata/metrics without forcing a DB migration for MVP. - [ ] The Agents UI has a readable advisor card and Storybook coverage. - [ ] Dogfooding produced screenshots and repro videos for the required scenarios. - [ ] Validation commands (`make lint`, targeted `make test`, Storybook tests, `make pre-commit`) passed before handoff. ## Suggested PR split 1. **PR 1 — Backend foundation** - `chatadvisor/` package - `chattool/advisor.go` - `chatloop` exclusive policy - chatd gating/prompt sync - backend tests 2. **PR 2 — Frontend + QA** - advisor renderer - stories/play assertions - dogfood artifacts and QA notes 3. **PR 3 — Optional follow-ups only if demanded by stakeholders** - separate advisor model override - persistent advisor billing/queryability - transient phase-stream UX </details> --- _Generated with [`mux`](https://github.com/coder/mux) • Model: `anthropic:claude-opus-4-7` • Thinking: `max`_
This commit is contained in:
Thomas Kosiewski
@@ -4315,6 +4315,8 @@ func (api *API) putChatAdvisorConfig(rw http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
publishChatConfigEvent(api.Logger, api.Pubsub, pubsub.ChatConfigEventAdvisorConfig, uuid.Nil)
|
||||
|
||||
rw.WriteHeader(http.StatusNoContent)
|
||||
}
|
||||
|
||||
|
||||
@@ -9,8 +9,9 @@ import (
|
||||
)
|
||||
|
||||
// ChatConfigEventChannel is the pubsub channel for chat config
|
||||
// changes (providers, model configs, user prompts). All replicas
|
||||
// subscribe to this channel to invalidate their local caches.
|
||||
// changes (providers, model configs, user prompts, advisor config).
|
||||
// All replicas subscribe to this channel to invalidate their local
|
||||
// caches.
|
||||
const ChatConfigEventChannel = "chat:config_change"
|
||||
|
||||
// HandleChatConfigEvent wraps a typed callback for ChatConfigEvent
|
||||
@@ -32,21 +33,24 @@ func HandleChatConfigEvent(cb func(ctx context.Context, payload ChatConfigEvent,
|
||||
}
|
||||
|
||||
// ChatConfigEvent is published when chat configuration changes
|
||||
// (provider CRUD, model config CRUD, or user prompt updates).
|
||||
// Subscribers use this to invalidate their local caches.
|
||||
// (provider CRUD, model config CRUD, user prompt updates, or advisor
|
||||
// config updates). Subscribers use this to invalidate their local
|
||||
// caches.
|
||||
type ChatConfigEvent struct {
|
||||
Kind ChatConfigEventKind `json:"kind"`
|
||||
// EntityID carries context for the invalidation:
|
||||
// - For providers: uuid.Nil (all providers are invalidated).
|
||||
// - For model configs: the specific config ID.
|
||||
// - For user prompts: the user ID.
|
||||
// - For advisor config: uuid.Nil (singleton site-config row).
|
||||
EntityID uuid.UUID `json:"entity_id"`
|
||||
}
|
||||
|
||||
type ChatConfigEventKind string
|
||||
|
||||
const (
|
||||
ChatConfigEventProviders ChatConfigEventKind = "providers"
|
||||
ChatConfigEventModelConfig ChatConfigEventKind = "model_config"
|
||||
ChatConfigEventUserPrompt ChatConfigEventKind = "user_prompt"
|
||||
ChatConfigEventProviders ChatConfigEventKind = "providers"
|
||||
ChatConfigEventModelConfig ChatConfigEventKind = "model_config"
|
||||
ChatConfigEventUserPrompt ChatConfigEventKind = "user_prompt"
|
||||
ChatConfigEventAdvisorConfig ChatConfigEventKind = "advisor_config"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,459 @@
|
||||
package chatd //nolint:testpackage // Accesses unexported advisor helpers.
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"charm.land/fantasy"
|
||||
fantasyopenai "charm.land/fantasy/providers/openai"
|
||||
"github.com/google/uuid"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"cdr.dev/slog/v3"
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatadvisor"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprovider"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chattest"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
// advisorOverrideStubStore stubs only the database methods that
|
||||
// resolveAdvisorModelOverride exercises. The prod code calls
|
||||
// GetEnabledChatModelConfigByID so the query joins chat_providers and
|
||||
// filters both enabled flags atomically; tests simulate that by returning
|
||||
// configs the stub treats as enabled.
|
||||
type advisorOverrideStubStore struct {
|
||||
database.Store
|
||||
|
||||
getEnabledChatModelConfigByID func(context.Context, uuid.UUID) (database.ChatModelConfig, error)
|
||||
}
|
||||
|
||||
func (s *advisorOverrideStubStore) GetEnabledChatModelConfigByID(
|
||||
ctx context.Context,
|
||||
id uuid.UUID,
|
||||
) (database.ChatModelConfig, error) {
|
||||
if s.getEnabledChatModelConfigByID == nil {
|
||||
return database.ChatModelConfig{}, xerrors.New("unexpected GetEnabledChatModelConfigByID call")
|
||||
}
|
||||
return s.getEnabledChatModelConfigByID(ctx, id)
|
||||
}
|
||||
|
||||
func newAdvisorTestServer(
|
||||
ctx context.Context,
|
||||
t *testing.T,
|
||||
store database.Store,
|
||||
) *Server {
|
||||
t.Helper()
|
||||
clock := quartz.NewMock(t)
|
||||
return &Server{
|
||||
db: store,
|
||||
configCache: newChatConfigCache(ctx, store, clock),
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveAdvisorModelOverride covers the early-return, each fallback
|
||||
// branch, and the success path. Prior tests only hit the ModelConfigID ==
|
||||
// uuid.Nil early return, so the override body never executed.
|
||||
func TestResolveAdvisorModelOverride(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
fallbackModel := &chattest.FakeModel{ProviderName: "stub", ModelName: "stub"}
|
||||
fallbackCallConfig := codersdk.ChatModelCallConfig{}
|
||||
logger := slog.Make()
|
||||
|
||||
t.Run("NilModelConfigReturnsFallback", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
// Panic if the cache is consulted; the early return must skip it.
|
||||
store := &advisorOverrideStubStore{}
|
||||
p := newAdvisorTestServer(ctx, t, store)
|
||||
|
||||
gotModel, gotCfg := p.resolveAdvisorModelOverride(
|
||||
ctx,
|
||||
database.Chat{},
|
||||
codersdk.AdvisorConfig{},
|
||||
fallbackModel,
|
||||
fallbackCallConfig,
|
||||
chatprovider.ProviderAPIKeys{},
|
||||
logger,
|
||||
)
|
||||
require.Equal(t, fallbackModel, gotModel)
|
||||
require.Equal(t, fallbackCallConfig, gotCfg)
|
||||
})
|
||||
|
||||
t.Run("ConfigLookupErrorReturnsFallback", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
store := &advisorOverrideStubStore{
|
||||
getEnabledChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) {
|
||||
return database.ChatModelConfig{}, xerrors.New("lookup failed")
|
||||
},
|
||||
}
|
||||
p := newAdvisorTestServer(ctx, t, store)
|
||||
|
||||
gotModel, gotCfg := p.resolveAdvisorModelOverride(
|
||||
ctx,
|
||||
database.Chat{},
|
||||
codersdk.AdvisorConfig{ModelConfigID: uuid.New()},
|
||||
fallbackModel,
|
||||
fallbackCallConfig,
|
||||
chatprovider.ProviderAPIKeys{OpenAI: "sk-test"},
|
||||
logger,
|
||||
)
|
||||
require.Equal(t, fallbackModel, gotModel)
|
||||
require.Equal(t, fallbackCallConfig, gotCfg)
|
||||
})
|
||||
|
||||
// Covers the sql.ErrNoRows branch separately from the generic-error
|
||||
// branch above. GetEnabledChatModelConfigByID returns ErrNoRows when
|
||||
// an admin disables the advisor model or its provider, and that case
|
||||
// has a distinct log message. Without this test, removing the
|
||||
// errors.Is(err, sql.ErrNoRows) check would still pass the sibling
|
||||
// test.
|
||||
t.Run("DisabledProviderReturnsFallback", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
store := &advisorOverrideStubStore{
|
||||
getEnabledChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) {
|
||||
return database.ChatModelConfig{}, sql.ErrNoRows
|
||||
},
|
||||
}
|
||||
p := newAdvisorTestServer(ctx, t, store)
|
||||
|
||||
gotModel, gotCfg := p.resolveAdvisorModelOverride(
|
||||
ctx,
|
||||
database.Chat{},
|
||||
codersdk.AdvisorConfig{ModelConfigID: uuid.New()},
|
||||
fallbackModel,
|
||||
fallbackCallConfig,
|
||||
chatprovider.ProviderAPIKeys{OpenAI: "sk-test"},
|
||||
logger,
|
||||
)
|
||||
require.Equal(t, fallbackModel, gotModel)
|
||||
require.Equal(t, fallbackCallConfig, gotCfg)
|
||||
})
|
||||
|
||||
t.Run("InvalidOptionsJSONReturnsFallback", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
configID := uuid.New()
|
||||
store := &advisorOverrideStubStore{
|
||||
getEnabledChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) {
|
||||
return database.ChatModelConfig{
|
||||
ID: configID,
|
||||
Provider: "openai",
|
||||
Model: "gpt-5.2",
|
||||
Enabled: true,
|
||||
CreatedAt: time.Unix(0, 0).UTC(),
|
||||
UpdatedAt: time.Unix(0, 0).UTC(),
|
||||
Options: []byte("not valid json"),
|
||||
DisplayName: "gpt-5.2",
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
p := newAdvisorTestServer(ctx, t, store)
|
||||
|
||||
gotModel, gotCfg := p.resolveAdvisorModelOverride(
|
||||
ctx,
|
||||
database.Chat{},
|
||||
codersdk.AdvisorConfig{ModelConfigID: configID},
|
||||
fallbackModel,
|
||||
fallbackCallConfig,
|
||||
chatprovider.ProviderAPIKeys{OpenAI: "sk-test"},
|
||||
logger,
|
||||
)
|
||||
require.Equal(t, fallbackModel, gotModel)
|
||||
require.Equal(t, fallbackCallConfig, gotCfg)
|
||||
})
|
||||
|
||||
t.Run("MissingProviderKeyReturnsFallback", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
configID := uuid.New()
|
||||
store := &advisorOverrideStubStore{
|
||||
getEnabledChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) {
|
||||
return database.ChatModelConfig{
|
||||
ID: configID,
|
||||
Provider: "openai",
|
||||
Model: "gpt-5.2",
|
||||
Enabled: true,
|
||||
CreatedAt: time.Unix(0, 0).UTC(),
|
||||
UpdatedAt: time.Unix(0, 0).UTC(),
|
||||
DisplayName: "gpt-5.2",
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
p := newAdvisorTestServer(ctx, t, store)
|
||||
|
||||
gotModel, gotCfg := p.resolveAdvisorModelOverride(
|
||||
ctx,
|
||||
database.Chat{},
|
||||
codersdk.AdvisorConfig{ModelConfigID: configID},
|
||||
fallbackModel,
|
||||
fallbackCallConfig,
|
||||
chatprovider.ProviderAPIKeys{},
|
||||
logger,
|
||||
)
|
||||
require.Equal(t, fallbackModel, gotModel)
|
||||
require.Equal(t, fallbackCallConfig, gotCfg)
|
||||
})
|
||||
|
||||
t.Run("SuccessReturnsOverrideModelAndConfig", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
configID := uuid.New()
|
||||
rawOptions, err := json.Marshal(codersdk.ChatModelCallConfig{
|
||||
Temperature: func() *float64 { v := 0.42; return &v }(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
store := &advisorOverrideStubStore{
|
||||
getEnabledChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) {
|
||||
return database.ChatModelConfig{
|
||||
ID: configID,
|
||||
Provider: "openai",
|
||||
Model: "gpt-5.2",
|
||||
Enabled: true,
|
||||
CreatedAt: time.Unix(0, 0).UTC(),
|
||||
UpdatedAt: time.Unix(0, 0).UTC(),
|
||||
Options: rawOptions,
|
||||
DisplayName: "gpt-5.2",
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
p := newAdvisorTestServer(ctx, t, store)
|
||||
|
||||
gotModel, gotCfg := p.resolveAdvisorModelOverride(
|
||||
ctx,
|
||||
database.Chat{},
|
||||
codersdk.AdvisorConfig{ModelConfigID: configID},
|
||||
fallbackModel,
|
||||
fallbackCallConfig,
|
||||
chatprovider.ProviderAPIKeys{OpenAI: "sk-test"},
|
||||
logger,
|
||||
)
|
||||
require.NotEqual(t, fantasy.LanguageModel(fallbackModel), gotModel,
|
||||
"success path must return the override model, not the fallback")
|
||||
require.NotNil(t, gotModel)
|
||||
require.Equal(t, "openai", gotModel.Provider())
|
||||
// Guard against ModelFromConfig silently ignoring the model field
|
||||
// and returning a default. The override is only useful if the
|
||||
// model name from the config row actually propagates.
|
||||
require.Equal(t, "gpt-5.2", gotModel.Model())
|
||||
require.NotNil(t, gotCfg.Temperature)
|
||||
require.InDelta(t, 0.42, *gotCfg.Temperature, 1e-9)
|
||||
})
|
||||
}
|
||||
|
||||
// TestStripAdvisorGuidanceBlock exercises the filter that keeps the advisor
|
||||
// from receiving the parent-facing advisor-guidance instruction in its nested
|
||||
// context. The block references a tool the advisor cannot use, so forwarding
|
||||
// it wastes context tokens and risks steering the advisor's reply.
|
||||
func TestStripAdvisorGuidanceBlock(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("RemovesGuidanceSystemMessage", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
msgs := []fantasy.Message{
|
||||
{
|
||||
Role: fantasy.MessageRoleSystem,
|
||||
Content: []fantasy.MessagePart{
|
||||
fantasy.TextPart{Text: "You are a helpful assistant."},
|
||||
},
|
||||
},
|
||||
{
|
||||
Role: fantasy.MessageRoleSystem,
|
||||
Content: []fantasy.MessagePart{
|
||||
fantasy.TextPart{Text: chatadvisor.ParentGuidanceBlock},
|
||||
},
|
||||
},
|
||||
{
|
||||
Role: fantasy.MessageRoleUser,
|
||||
Content: []fantasy.MessagePart{
|
||||
fantasy.TextPart{Text: "Help me plan."},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
filtered := stripAdvisorGuidanceBlock(msgs)
|
||||
require.Len(t, filtered, 2)
|
||||
for _, msg := range filtered {
|
||||
for _, part := range msg.Content {
|
||||
if text, ok := part.(fantasy.TextPart); ok {
|
||||
require.NotEqual(t, chatadvisor.ParentGuidanceBlock, text.Text,
|
||||
"guidance block must not survive the filter")
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("LeavesOtherSystemMessagesIntact", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
msgs := []fantasy.Message{
|
||||
{
|
||||
Role: fantasy.MessageRoleSystem,
|
||||
Content: []fantasy.MessagePart{
|
||||
fantasy.TextPart{Text: "instruction file"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Role: fantasy.MessageRoleUser,
|
||||
Content: []fantasy.MessagePart{
|
||||
fantasy.TextPart{Text: "hi"},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
filtered := stripAdvisorGuidanceBlock(msgs)
|
||||
require.Len(t, filtered, 2)
|
||||
})
|
||||
|
||||
t.Run("IgnoresNonSystemRoleWithMatchingText", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
// A user message echoing the guidance block must not be stripped:
|
||||
// the filter only targets the system-role injection.
|
||||
msgs := []fantasy.Message{
|
||||
{
|
||||
Role: fantasy.MessageRoleUser,
|
||||
Content: []fantasy.MessagePart{
|
||||
fantasy.TextPart{Text: chatadvisor.ParentGuidanceBlock},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
filtered := stripAdvisorGuidanceBlock(msgs)
|
||||
require.Len(t, filtered, 1)
|
||||
})
|
||||
}
|
||||
|
||||
// TestNewAdvisorRuntime covers the three defensive branches in
|
||||
// newAdvisorRuntime that gate whether the runtime is created and with what
|
||||
// bounds. Without this coverage a regression in any branch ships silently.
|
||||
func TestNewAdvisorRuntime(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
logger := slog.Make()
|
||||
fallbackModel := &chattest.FakeModel{ProviderName: "openai", ModelName: "gpt-4"}
|
||||
fallbackCallConfig := codersdk.ChatModelCallConfig{}
|
||||
|
||||
t.Run("ZeroMaxUsesDefaultsToMaxChatSteps", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
store := &advisorOverrideStubStore{}
|
||||
p := newAdvisorTestServer(ctx, t, store)
|
||||
|
||||
rt := p.newAdvisorRuntime(
|
||||
ctx,
|
||||
database.Chat{},
|
||||
codersdk.AdvisorConfig{
|
||||
Enabled: true,
|
||||
MaxUsesPerRun: 0,
|
||||
MaxOutputTokens: 16384,
|
||||
},
|
||||
fallbackModel,
|
||||
fallbackCallConfig,
|
||||
chatprovider.ProviderAPIKeys{},
|
||||
logger,
|
||||
)
|
||||
require.NotNil(t, rt, "zero max uses must default rather than bail out")
|
||||
require.Equal(t, maxChatSteps, rt.RemainingUses(),
|
||||
"zero max uses must be replaced with maxChatSteps")
|
||||
})
|
||||
|
||||
t.Run("NegativeMaxUsesReturnsNil", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
store := &advisorOverrideStubStore{}
|
||||
p := newAdvisorTestServer(ctx, t, store)
|
||||
|
||||
rt := p.newAdvisorRuntime(
|
||||
ctx,
|
||||
database.Chat{},
|
||||
codersdk.AdvisorConfig{
|
||||
Enabled: true,
|
||||
MaxUsesPerRun: -1,
|
||||
MaxOutputTokens: 16384,
|
||||
},
|
||||
fallbackModel,
|
||||
fallbackCallConfig,
|
||||
chatprovider.ProviderAPIKeys{},
|
||||
logger,
|
||||
)
|
||||
require.Nil(t, rt, "negative max uses must disable the advisor")
|
||||
})
|
||||
|
||||
t.Run("ZeroMaxOutputTokensDefaults", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
store := &advisorOverrideStubStore{}
|
||||
p := newAdvisorTestServer(ctx, t, store)
|
||||
|
||||
rt := p.newAdvisorRuntime(
|
||||
ctx,
|
||||
database.Chat{},
|
||||
codersdk.AdvisorConfig{
|
||||
Enabled: true,
|
||||
MaxUsesPerRun: 3,
|
||||
MaxOutputTokens: 0,
|
||||
},
|
||||
fallbackModel,
|
||||
fallbackCallConfig,
|
||||
chatprovider.ProviderAPIKeys{},
|
||||
logger,
|
||||
)
|
||||
require.NotNil(t, rt,
|
||||
"zero max output tokens must default to defaultAdvisorMaxOutputTokens, not disable the advisor")
|
||||
require.Equal(t, 3, rt.RemainingUses())
|
||||
require.Equal(t, int64(defaultAdvisorMaxOutputTokens), rt.MaxOutputTokens(),
|
||||
"zero max output tokens must be replaced with defaultAdvisorMaxOutputTokens")
|
||||
})
|
||||
|
||||
// Guards the wiring from AdvisorConfig.ReasoningEffort through
|
||||
// newAdvisorRuntime to ApplyReasoningEffortToOptions. A field swap,
|
||||
// typo, or accidental deletion of the apply call would otherwise
|
||||
// ship silently because chatprovider_test only covers the helper in
|
||||
// isolation.
|
||||
t.Run("ReasoningEffortReachesProviderOptions", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
store := &advisorOverrideStubStore{}
|
||||
p := newAdvisorTestServer(ctx, t, store)
|
||||
|
||||
openAIModel := &chattest.FakeModel{
|
||||
ProviderName: fantasyopenai.Name,
|
||||
ModelName: "gpt-4",
|
||||
}
|
||||
|
||||
rt := p.newAdvisorRuntime(
|
||||
ctx,
|
||||
database.Chat{},
|
||||
codersdk.AdvisorConfig{
|
||||
Enabled: true,
|
||||
MaxUsesPerRun: 3,
|
||||
MaxOutputTokens: 16384,
|
||||
ReasoningEffort: "high",
|
||||
},
|
||||
openAIModel,
|
||||
fallbackCallConfig,
|
||||
chatprovider.ProviderAPIKeys{},
|
||||
logger,
|
||||
)
|
||||
require.NotNil(t, rt)
|
||||
|
||||
providerOptions := rt.ProviderOptions()
|
||||
require.NotNil(t, providerOptions,
|
||||
"advisor runtime must seed provider options when reasoning effort is set")
|
||||
opts, ok := providerOptions[fantasyopenai.Name].(*fantasyopenai.ResponsesProviderOptions)
|
||||
require.True(t, ok,
|
||||
"expected *ResponsesProviderOptions for Responses model, got %T",
|
||||
providerOptions[fantasyopenai.Name])
|
||||
require.NotNil(t, opts.ReasoningEffort,
|
||||
"ReasoningEffort from AdvisorConfig must reach the provider options")
|
||||
require.Equal(t, fantasyopenai.ReasoningEffortHigh, *opts.ReasoningEffort)
|
||||
})
|
||||
}
|
||||
@@ -121,6 +121,28 @@ func (rt *Runtime) RemainingUses() int {
|
||||
return int(remaining)
|
||||
}
|
||||
|
||||
// MaxOutputTokens reports the resolved output-token cap applied to each
|
||||
// advisor call. NewRuntime validates that this value is positive and that
|
||||
// it matches ModelConfig.MaxOutputTokens when both are set, so the
|
||||
// accessor always returns the value the runtime will actually send.
|
||||
func (rt *Runtime) MaxOutputTokens() int64 {
|
||||
if rt == nil {
|
||||
return 0
|
||||
}
|
||||
return rt.cfg.MaxOutputTokens
|
||||
}
|
||||
|
||||
// ProviderOptions reports the resolved provider options applied to each
|
||||
// advisor call. NewRuntime clones the supplied options so the returned
|
||||
// map reflects what nested calls will actually receive; callers must not
|
||||
// mutate the map or its entries.
|
||||
func (rt *Runtime) ProviderOptions() fantasy.ProviderOptions {
|
||||
if rt == nil {
|
||||
return nil
|
||||
}
|
||||
return rt.cfg.ProviderOptions
|
||||
}
|
||||
|
||||
func (rt *Runtime) tryAcquire() bool {
|
||||
for {
|
||||
used := rt.used.Load()
|
||||
|
||||
@@ -10,6 +10,11 @@ import (
|
||||
"charm.land/fantasy"
|
||||
)
|
||||
|
||||
// ToolName is the identifier the advisor tool registers under. The parent
|
||||
// agent's exclusive-tool policy and the advisor-guidance block both reference
|
||||
// this name, so keeping them synchronized requires a single source of truth.
|
||||
const ToolName = "advisor"
|
||||
|
||||
// advisorQuestionMaxRunes caps the parent agent's question at a length
|
||||
// that leaves room in the advisor prompt for system preamble and recent
|
||||
// conversation context.
|
||||
@@ -26,7 +31,7 @@ type ToolOptions struct {
|
||||
// context, runs without tools, and is limited to a single model step.
|
||||
func Tool(opts ToolOptions) fantasy.AgentTool {
|
||||
return fantasy.NewAgentTool(
|
||||
"advisor",
|
||||
ToolName,
|
||||
"Ask a separate advisor pass for strategic guidance about planning, architecture, tradeoffs, or debugging strategy. Provide a brief question. The advisor sees recent conversation context, runs without tools for a single step, and responds to the parent agent rather than the end user.",
|
||||
func(ctx context.Context, args AdvisorArgs, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
if opts.Runtime == nil {
|
||||
|
||||
+293
-9
@@ -36,6 +36,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/util/xjson"
|
||||
"github.com/coder/coder/v2/coderd/webpush"
|
||||
"github.com/coder/coder/v2/coderd/workspacestats"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatadvisor"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatcost"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatdebug"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chaterror"
|
||||
@@ -118,6 +119,12 @@ const (
|
||||
DefaultMaxChatsPerAcquire int32 = 10
|
||||
|
||||
defaultSubagentInstruction = "You are running as a delegated sub-agent chat. Complete the delegated task and provide clear, concise assistant responses for the parent agent."
|
||||
|
||||
// defaultAdvisorMaxOutputTokens caps the nested advisor response
|
||||
// when the admin config omits the field (or sets it to <= 0).
|
||||
// It is intentionally generous relative to the advisor's concise
|
||||
// guidance remit so short plans are not truncated mid-reasoning.
|
||||
defaultAdvisorMaxOutputTokens = 16384
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -225,6 +232,192 @@ func (p *Server) chatTemplateAllowlist() map[uuid.UUID]bool {
|
||||
return m
|
||||
}
|
||||
|
||||
func (p *Server) loadAdvisorConfig(ctx context.Context, logger slog.Logger) codersdk.AdvisorConfig {
|
||||
cfg, err := p.configCache.AdvisorConfig(ctx)
|
||||
if err != nil {
|
||||
logger.Warn(ctx, "failed to load advisor config", slog.Error(err))
|
||||
return codersdk.AdvisorConfig{}
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
// stripAdvisorGuidanceBlock removes any system message whose text content
|
||||
// matches chatadvisor.ParentGuidanceBlock after whitespace normalization.
|
||||
// The block is meant for the parent agent (it advertises the advisor tool)
|
||||
// and would waste context tokens if forwarded to the advisor's nested run.
|
||||
func stripAdvisorGuidanceBlock(msgs []fantasy.Message) []fantasy.Message {
|
||||
filtered := msgs[:0]
|
||||
for _, msg := range msgs {
|
||||
if msg.Role == fantasy.MessageRoleSystem && isAdvisorGuidanceMessage(msg) {
|
||||
continue
|
||||
}
|
||||
filtered = append(filtered, msg)
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
func isAdvisorGuidanceMessage(msg fantasy.Message) bool {
|
||||
if len(msg.Content) != 1 {
|
||||
return false
|
||||
}
|
||||
text, ok := msg.Content[0].(fantasy.TextPart)
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
return strings.TrimSpace(text.Text) == strings.TrimSpace(chatadvisor.ParentGuidanceBlock)
|
||||
}
|
||||
|
||||
func (p *Server) resolveAdvisorModelOverride(
|
||||
ctx context.Context,
|
||||
chat database.Chat,
|
||||
advisorCfg codersdk.AdvisorConfig,
|
||||
fallbackModel fantasy.LanguageModel,
|
||||
fallbackCallConfig codersdk.ChatModelCallConfig,
|
||||
providerKeys chatprovider.ProviderAPIKeys,
|
||||
logger slog.Logger,
|
||||
) (fantasy.LanguageModel, codersdk.ChatModelCallConfig) {
|
||||
if advisorCfg.ModelConfigID == uuid.Nil {
|
||||
return fallbackModel, fallbackCallConfig
|
||||
}
|
||||
|
||||
// GetEnabledChatModelConfigByID joins on chat_providers.enabled = TRUE
|
||||
// and chat_model_configs.enabled = TRUE, so it returns sql.ErrNoRows
|
||||
// the moment an admin disables either the model config or its provider.
|
||||
// Using the cached ModelConfigByID here would keep resolving an override
|
||||
// whose provider was just disabled, and an env or central fallback key
|
||||
// would let ModelFromConfig succeed, silently routing advisor prompts
|
||||
// to a provider the admin expects to be off.
|
||||
overrideConfig, err := p.db.GetEnabledChatModelConfigByID(
|
||||
ctx,
|
||||
advisorCfg.ModelConfigID,
|
||||
)
|
||||
if err != nil {
|
||||
if xerrors.Is(err, sql.ErrNoRows) {
|
||||
logger.Warn(
|
||||
ctx,
|
||||
"advisor model config is disabled or unavailable, continuing with chat model",
|
||||
slog.F("model_config_id", advisorCfg.ModelConfigID),
|
||||
)
|
||||
return fallbackModel, fallbackCallConfig
|
||||
}
|
||||
logger.Warn(
|
||||
ctx,
|
||||
"failed to resolve advisor model config, continuing with chat model",
|
||||
slog.F("model_config_id", advisorCfg.ModelConfigID),
|
||||
slog.Error(err),
|
||||
)
|
||||
return fallbackModel, fallbackCallConfig
|
||||
}
|
||||
|
||||
overrideCallConfig := codersdk.ChatModelCallConfig{}
|
||||
if len(overrideConfig.Options) > 0 {
|
||||
if err := json.Unmarshal(overrideConfig.Options, &overrideCallConfig); err != nil {
|
||||
logger.Warn(
|
||||
ctx,
|
||||
"failed to parse advisor model config, continuing with chat model",
|
||||
slog.F("model_config_id", advisorCfg.ModelConfigID),
|
||||
slog.Error(err),
|
||||
)
|
||||
return fallbackModel, fallbackCallConfig
|
||||
}
|
||||
}
|
||||
|
||||
overrideModel, err := chatprovider.ModelFromConfig(
|
||||
overrideConfig.Provider,
|
||||
overrideConfig.Model,
|
||||
providerKeys,
|
||||
chatprovider.UserAgent(),
|
||||
chatprovider.CoderHeaders(chat),
|
||||
nil,
|
||||
)
|
||||
if err != nil {
|
||||
logger.Warn(
|
||||
ctx,
|
||||
"failed to create advisor override model, continuing with chat model",
|
||||
slog.F("model_config_id", advisorCfg.ModelConfigID),
|
||||
slog.Error(err),
|
||||
)
|
||||
return fallbackModel, fallbackCallConfig
|
||||
}
|
||||
|
||||
return overrideModel, overrideCallConfig
|
||||
}
|
||||
|
||||
func (p *Server) newAdvisorRuntime(
|
||||
ctx context.Context,
|
||||
chat database.Chat,
|
||||
advisorCfg codersdk.AdvisorConfig,
|
||||
fallbackModel fantasy.LanguageModel,
|
||||
fallbackCallConfig codersdk.ChatModelCallConfig,
|
||||
providerKeys chatprovider.ProviderAPIKeys,
|
||||
logger slog.Logger,
|
||||
) *chatadvisor.Runtime {
|
||||
advisorModel, advisorCallConfig := p.resolveAdvisorModelOverride(
|
||||
ctx,
|
||||
chat,
|
||||
advisorCfg,
|
||||
fallbackModel,
|
||||
fallbackCallConfig,
|
||||
providerKeys,
|
||||
logger,
|
||||
)
|
||||
|
||||
maxUsesPerRun := advisorCfg.MaxUsesPerRun
|
||||
switch {
|
||||
case maxUsesPerRun == 0:
|
||||
// Advisor config treats 0 as unlimited, but the runtime
|
||||
// requires a positive bound. maxChatSteps is the
|
||||
// effective upper bound because advisor can run at most
|
||||
// once per loop step.
|
||||
maxUsesPerRun = maxChatSteps
|
||||
case maxUsesPerRun < 0:
|
||||
logger.Warn(
|
||||
ctx,
|
||||
"invalid advisor max uses per run, continuing without advisor",
|
||||
slog.F("max_uses_per_run", maxUsesPerRun),
|
||||
)
|
||||
return nil
|
||||
}
|
||||
|
||||
maxOutputTokens := advisorCfg.MaxOutputTokens
|
||||
if maxOutputTokens <= 0 {
|
||||
maxOutputTokens = defaultAdvisorMaxOutputTokens
|
||||
}
|
||||
|
||||
advisorCallConfig.MaxOutputTokens = ptr.Ref(maxOutputTokens)
|
||||
providerOptions := chatprovider.ProviderOptionsFromChatModelConfig(
|
||||
advisorModel,
|
||||
advisorCallConfig.ProviderOptions,
|
||||
)
|
||||
// ProviderOptionsFromChatModelConfig returns nil when the model config
|
||||
// has no provider_options block, so the helper seeds a minimal entry
|
||||
// for the advisor model's provider before applying reasoning_effort.
|
||||
// This keeps the per-provider dispatch in chatprovider so adding a new
|
||||
// provider there propagates here automatically.
|
||||
providerOptions = chatprovider.ApplyReasoningEffortToOptions(
|
||||
providerOptions,
|
||||
advisorModel,
|
||||
advisorCfg.ReasoningEffort,
|
||||
)
|
||||
|
||||
rt, err := chatadvisor.NewRuntime(chatadvisor.RuntimeConfig{
|
||||
Model: advisorModel,
|
||||
ModelConfig: advisorCallConfig,
|
||||
ProviderOptions: providerOptions,
|
||||
MaxUsesPerRun: maxUsesPerRun,
|
||||
MaxOutputTokens: maxOutputTokens,
|
||||
})
|
||||
if err != nil {
|
||||
logger.Warn(
|
||||
ctx,
|
||||
"failed to create advisor runtime, continuing without advisor",
|
||||
slog.Error(err),
|
||||
)
|
||||
return nil
|
||||
}
|
||||
return rt
|
||||
}
|
||||
|
||||
// cachedWorkspaceMCPTools stores workspace MCP tools discovered
|
||||
// from a workspace agent, keyed by the agent ID that provided them.
|
||||
type cachedWorkspaceMCPTools struct {
|
||||
@@ -3754,6 +3947,8 @@ func New(cfg Config) *Server {
|
||||
p.configCache.InvalidateModelConfig(ev.EntityID)
|
||||
case coderdpubsub.ChatConfigEventUserPrompt:
|
||||
p.configCache.InvalidateUserPrompt(ev.EntityID)
|
||||
case coderdpubsub.ChatConfigEventAdvisorConfig:
|
||||
p.configCache.InvalidateAdvisorConfig()
|
||||
}
|
||||
}),
|
||||
)
|
||||
@@ -6034,6 +6229,48 @@ func (p *Server) runChat(
|
||||
}
|
||||
planModeInstructions := p.loadPlanModeInstructions(ctx, currentPlanMode, logger)
|
||||
|
||||
advisorCfg := p.loadAdvisorConfig(ctx, logger)
|
||||
|
||||
var advisorRuntime *chatadvisor.Runtime
|
||||
// Plan mode filters the advisor tool out of the turn's tool set via
|
||||
// filterToolsForTurn, so enabling the runtime there would inject
|
||||
// guidance and enforce advisor exclusivity for a tool the model
|
||||
// cannot actually call. Explore chats (root or subagent) run under
|
||||
// allowedExploreToolNames, whose policy does not include advisor, so
|
||||
// registering the runtime there would inject guidance for a tool
|
||||
// that is never exposed to the model.
|
||||
if advisorCfg.Enabled && isRootChat && !isPlanModeTurn && !isExploreSubagent {
|
||||
advisorRuntime = p.newAdvisorRuntime(
|
||||
ctx,
|
||||
chat,
|
||||
advisorCfg,
|
||||
model,
|
||||
callConfig,
|
||||
providerKeys,
|
||||
logger,
|
||||
)
|
||||
}
|
||||
|
||||
var advisorPromptSnapshot []fantasy.Message
|
||||
// setAdvisorPromptSnapshot captures the final prompt state the outer
|
||||
// model sees so the advisor tool can forward it as nested context.
|
||||
// It is invoked at four lifecycle points (after initial system-prompt
|
||||
// assembly, inside PrepareMessages before and after instruction
|
||||
// injection, and after ReloadMessages rebuilds the prompt) because
|
||||
// the prompt mutates at each of them and the advisor must snapshot
|
||||
// the post-mutation state. Removing any of those calls would leave
|
||||
// the advisor with a stale view of the conversation.
|
||||
//
|
||||
// The no-op guard keeps the common disabled/filtered paths (advisor
|
||||
// off, plan mode, explore, child chats) from paying an O(n) prompt
|
||||
// clone per step for a snapshot that is never consumed.
|
||||
setAdvisorPromptSnapshot := func(msgs []fantasy.Message) {
|
||||
if advisorRuntime == nil {
|
||||
return
|
||||
}
|
||||
advisorPromptSnapshot = slices.Clone(msgs)
|
||||
}
|
||||
|
||||
chainInfo := resolveChainMode(messages)
|
||||
result.PushSummaryModel = model
|
||||
result.ProviderKeys = providerKeys
|
||||
@@ -6336,6 +6573,10 @@ func (p *Server) runChat(
|
||||
isRootChat: isRootChat,
|
||||
},
|
||||
)
|
||||
// Inject advisor guidance when the advisor runtime is available.
|
||||
if advisorRuntime != nil {
|
||||
prompt = chatprompt.InsertSystem(prompt, chatadvisor.ParentGuidanceBlock)
|
||||
}
|
||||
if mcpCleanup != nil {
|
||||
defer mcpCleanup()
|
||||
}
|
||||
@@ -6352,6 +6593,7 @@ func (p *Server) runChat(
|
||||
|
||||
instructionInjected := instruction != ""
|
||||
prompt = renderPlanPathPrompt(prompt, resolvePlanPathBlock(ctx))
|
||||
setAdvisorPromptSnapshot(prompt)
|
||||
// Use the model config's context_limit as a fallback when the LLM
|
||||
// provider doesn't include context_limit in its response metadata
|
||||
// (which is the common case).
|
||||
@@ -6742,6 +6984,24 @@ func (p *Server) runChat(
|
||||
chattool.ReadSkillFile(skillOpts),
|
||||
)
|
||||
}
|
||||
if advisorRuntime != nil {
|
||||
tools = append(tools, chatadvisor.Tool(chatadvisor.ToolOptions{
|
||||
Runtime: advisorRuntime,
|
||||
GetConversationSnapshot: func() []fantasy.Message {
|
||||
// The outer prompt contains ParentGuidanceBlock, which
|
||||
// tells the parent when to call the advisor tool. That
|
||||
// instruction is meaningless (and slightly confusing)
|
||||
// when forwarded to the advisor, whose nested run has
|
||||
// no tools. Strip it before handing the snapshot over.
|
||||
return stripAdvisorGuidanceBlock(slices.Clone(advisorPromptSnapshot))
|
||||
},
|
||||
}))
|
||||
}
|
||||
|
||||
var exclusiveToolNames map[string]bool
|
||||
if advisorRuntime != nil {
|
||||
exclusiveToolNames = map[string]bool{chatadvisor.ToolName: true}
|
||||
}
|
||||
|
||||
// Record builtin tool names before appending MCP tools
|
||||
// so the metrics layer can differentiate between built-in and MCP tools.
|
||||
@@ -6905,15 +7165,16 @@ func (p *Server) runChat(
|
||||
}()
|
||||
|
||||
loopErr = chatloop.Run(ctx, chatloop.RunOptions{
|
||||
Model: model,
|
||||
Messages: prompt,
|
||||
Tools: tools,
|
||||
ActiveTools: activeToolNames,
|
||||
StopAfterTools: stopAfterBehaviorTools(currentPlanMode, chat.Mode, chat.ParentChatID),
|
||||
MaxSteps: maxChatSteps,
|
||||
Metrics: p.metrics,
|
||||
Logger: loopLogger,
|
||||
BuiltinToolNames: builtinToolNames,
|
||||
Model: model,
|
||||
Messages: prompt,
|
||||
Tools: tools,
|
||||
ActiveTools: activeToolNames,
|
||||
StopAfterTools: stopAfterBehaviorTools(currentPlanMode, chat.Mode, chat.ParentChatID),
|
||||
MaxSteps: maxChatSteps,
|
||||
Metrics: p.metrics,
|
||||
Logger: loopLogger,
|
||||
BuiltinToolNames: builtinToolNames,
|
||||
ExclusiveToolNames: exclusiveToolNames,
|
||||
|
||||
ModelConfig: callConfig,
|
||||
ProviderOptions: providerOptions,
|
||||
@@ -6988,7 +7249,19 @@ func (p *Server) runChat(
|
||||
isRootChat: isRootChat,
|
||||
},
|
||||
)
|
||||
// Re-inject advisor guidance after rebuilding system
|
||||
// blocks so compaction/reload preserves the same
|
||||
// system-message ordering as the initial prompt path.
|
||||
if advisorRuntime != nil {
|
||||
reloadedPrompt = chatprompt.InsertSystem(reloadedPrompt, chatadvisor.ParentGuidanceBlock)
|
||||
}
|
||||
reloadedPrompt = renderPlanPathPrompt(reloadedPrompt, resolvePlanPathBlock(reloadCtx))
|
||||
// Snapshot the full reloaded prompt before chain-mode
|
||||
// filtering so the advisor runs with complete
|
||||
// assistant/tool context. The nested advisor call
|
||||
// clears previous_response_id, so provider-side
|
||||
// history is unavailable.
|
||||
setAdvisorPromptSnapshot(reloadedPrompt)
|
||||
if chainModeActive {
|
||||
reloadedPrompt = filterPromptForChainMode(
|
||||
reloadedPrompt,
|
||||
@@ -7001,6 +7274,14 @@ func (p *Server) runChat(
|
||||
chainModeActive = false
|
||||
},
|
||||
PrepareMessages: func(msgs []fantasy.Message) []fantasy.Message {
|
||||
// Skip the snapshot update when chain mode is active;
|
||||
// the chatloop passes in the chain-filtered prompt
|
||||
// (system plus trailing user messages) and the advisor
|
||||
// needs the full pre-chain history captured at the
|
||||
// initial-prompt and ReloadMessages sites.
|
||||
if !chainModeActive {
|
||||
setAdvisorPromptSnapshot(msgs)
|
||||
}
|
||||
if instructionInjected || instruction == "" {
|
||||
return nil
|
||||
}
|
||||
@@ -7009,6 +7290,9 @@ func (p *Server) runChat(
|
||||
if skillIndex := chattool.FormatSkillIndex(skills); skillIndex != "" {
|
||||
result = chatprompt.InsertSystem(result, skillIndex)
|
||||
}
|
||||
if !chainModeActive {
|
||||
setAdvisorPromptSnapshot(result)
|
||||
}
|
||||
return result
|
||||
},
|
||||
OnRetry: func(
|
||||
|
||||
@@ -44,6 +44,7 @@ import (
|
||||
"github.com/coder/coder/v2/coderd/util/slice"
|
||||
"github.com/coder/coder/v2/coderd/workspacestats"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatadvisor"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chatprompt"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chattest"
|
||||
"github.com/coder/coder/v2/coderd/x/chatd/chattool"
|
||||
@@ -8852,3 +8853,800 @@ func TestAcquireChatsSkipsArchivedPendingChat(t *testing.T) {
|
||||
require.Len(t, acquired, 1, "only the non-archived chat should be acquired")
|
||||
require.Equal(t, activeChat.ID, acquired[0].ID)
|
||||
}
|
||||
|
||||
func TestAdvisorGating_Disabled(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
var toolsMu sync.Mutex
|
||||
var capturedTools []string
|
||||
var capturedMessages []chattest.OpenAIMessage
|
||||
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
|
||||
names := make([]string, 0, len(req.Tools))
|
||||
for _, tool := range req.Tools {
|
||||
names = append(names, tool.Function.Name)
|
||||
}
|
||||
toolsMu.Lock()
|
||||
capturedTools = names
|
||||
capturedMessages = append([]chattest.OpenAIMessage(nil), req.Messages...)
|
||||
toolsMu.Unlock()
|
||||
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("advisor is not available")...,
|
||||
)
|
||||
})
|
||||
|
||||
user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
|
||||
seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{
|
||||
Enabled: false,
|
||||
MaxUsesPerRun: 3,
|
||||
MaxOutputTokens: 16384,
|
||||
})
|
||||
server := newActiveTestServer(t, db, ps)
|
||||
|
||||
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
Title: "advisor-disabled",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText("hello"),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
return got.Status == database.ChatStatusWaiting ||
|
||||
got.Status == database.ChatStatusError
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
toolsMu.Lock()
|
||||
tools := append([]string(nil), capturedTools...)
|
||||
messages := append([]chattest.OpenAIMessage(nil), capturedMessages...)
|
||||
toolsMu.Unlock()
|
||||
|
||||
require.NotEmpty(t, messages, "expected a streamed LLM request")
|
||||
require.NotContains(t, tools, "advisor",
|
||||
"advisor tool should not be registered when disabled")
|
||||
for _, msg := range messages {
|
||||
require.NotContains(t, msg.Content, chatadvisor.ParentGuidanceBlock,
|
||||
"advisor guidance should not be injected when disabled")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdvisorGating_RootChat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
var streamedCallCount atomic.Int32
|
||||
var streamedCallsMu sync.Mutex
|
||||
var firstCallTools []string
|
||||
var firstCallMessages []chattest.OpenAIMessage
|
||||
var secondCallMessages []chattest.OpenAIMessage
|
||||
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
|
||||
switch streamedCallCount.Add(1) {
|
||||
case 1:
|
||||
names := make([]string, 0, len(req.Tools))
|
||||
for _, tool := range req.Tools {
|
||||
names = append(names, tool.Function.Name)
|
||||
}
|
||||
streamedCallsMu.Lock()
|
||||
firstCallTools = names
|
||||
firstCallMessages = append([]chattest.OpenAIMessage(nil), req.Messages...)
|
||||
streamedCallsMu.Unlock()
|
||||
|
||||
advisorChunk := chattest.OpenAIToolCallChunk(
|
||||
"advisor",
|
||||
`{"question":"help me plan"}`,
|
||||
)
|
||||
readChunk := chattest.OpenAIToolCallChunk(
|
||||
"read_file",
|
||||
`{"path":"/tmp/test.txt"}`,
|
||||
)
|
||||
mergedChunk := advisorChunk
|
||||
readCall := readChunk.Choices[0].ToolCalls[0]
|
||||
readCall.Index = 1
|
||||
mergedChunk.Choices[0].ToolCalls = append(
|
||||
mergedChunk.Choices[0].ToolCalls,
|
||||
readCall,
|
||||
)
|
||||
return chattest.OpenAIStreamingResponse(mergedChunk)
|
||||
case 2:
|
||||
streamedCallsMu.Lock()
|
||||
secondCallMessages = append([]chattest.OpenAIMessage(nil), req.Messages...)
|
||||
streamedCallsMu.Unlock()
|
||||
}
|
||||
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("done")...,
|
||||
)
|
||||
})
|
||||
|
||||
user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
|
||||
seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{
|
||||
Enabled: true,
|
||||
MaxUsesPerRun: 3,
|
||||
MaxOutputTokens: 16384,
|
||||
})
|
||||
server := newActiveTestServer(t, db, ps)
|
||||
|
||||
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
Title: "advisor-root",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText("help me plan this"),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
if got.Status != database.ChatStatusWaiting &&
|
||||
got.Status != database.ChatStatusError {
|
||||
return false
|
||||
}
|
||||
return streamedCallCount.Load() >= 2
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
streamedCallsMu.Lock()
|
||||
tools := append([]string(nil), firstCallTools...)
|
||||
messages := append([]chattest.OpenAIMessage(nil), firstCallMessages...)
|
||||
secondMessages := append([]chattest.OpenAIMessage(nil), secondCallMessages...)
|
||||
streamedCallsMu.Unlock()
|
||||
|
||||
// Exactly two streamed LLM calls are expected: the first that
|
||||
// returned the mixed advisor + read_file batch, and the second
|
||||
// that received the exclusive-policy rejection. A third call
|
||||
// would indicate that either tool had slipped past the exclusive
|
||||
// policy; the >= 2 wait would have missed that regression.
|
||||
require.Equal(t, int32(2), streamedCallCount.Load(),
|
||||
"exclusive policy must block execution of both tools; no third call expected")
|
||||
require.NotEmpty(t, messages, "expected a first streamed LLM request")
|
||||
require.NotEmpty(t, secondMessages, "expected a second streamed LLM request")
|
||||
require.Contains(t, tools, "advisor",
|
||||
"advisor tool should be registered for root chats when enabled")
|
||||
|
||||
var hasGuidance bool
|
||||
for _, msg := range messages {
|
||||
if strings.Contains(msg.Content, chatadvisor.ParentGuidanceBlock) {
|
||||
hasGuidance = true
|
||||
break
|
||||
}
|
||||
}
|
||||
require.True(t, hasGuidance,
|
||||
"root chat should contain advisor guidance in the prompt")
|
||||
|
||||
var hasExclusiveAdvisorError bool
|
||||
var hasSkippedToolError bool
|
||||
for _, msg := range secondMessages {
|
||||
if strings.Contains(msg.Content, "advisor must be called alone") {
|
||||
hasExclusiveAdvisorError = true
|
||||
}
|
||||
if strings.Contains(msg.Content, "this tool was skipped because advisor must run alone") {
|
||||
hasSkippedToolError = true
|
||||
}
|
||||
}
|
||||
require.True(t, hasExclusiveAdvisorError,
|
||||
"mixed advisor batches should surface the exclusive advisor error")
|
||||
require.True(t, hasSkippedToolError,
|
||||
"mixed advisor batches should skip sibling tools with an explanatory error")
|
||||
}
|
||||
|
||||
// TestAdvisorHappyPath_RootChat walks the advisor tool end-to-end:
|
||||
// parent calls advisor alone, the nested advisor call produces text, and
|
||||
// the structured result flows back into the parent conversation. The
|
||||
// exclusive-policy test above only proves the rejection path; this test
|
||||
// covers the glue from chatd wiring -> chatadvisor.Tool -> Runtime.Run ->
|
||||
// nested model call -> structured result back to the outer model.
|
||||
func TestAdvisorHappyPath_RootChat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
const advisorReply = "break the problem into smaller pieces first"
|
||||
|
||||
var (
|
||||
streamedCallCount atomic.Int32
|
||||
streamedCallsMu sync.Mutex
|
||||
advisorCallSeen atomic.Bool
|
||||
advisorMessages []chattest.OpenAIMessage
|
||||
finalCallMessages []chattest.OpenAIMessage
|
||||
)
|
||||
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
|
||||
switch streamedCallCount.Add(1) {
|
||||
case 1:
|
||||
// Parent turn 1: call advisor solo.
|
||||
return chattest.OpenAIStreamingResponse(chattest.OpenAIToolCallChunk(
|
||||
"advisor",
|
||||
`{"question":"how should I approach this refactor?"}`,
|
||||
))
|
||||
case 2:
|
||||
// Nested advisor turn. The nested call has no tools because
|
||||
// chatadvisor.RunAdvisor runs with MaxSteps=1 and no tool
|
||||
// set.
|
||||
require.Empty(t, req.Tools,
|
||||
"advisor's nested call must run without tools")
|
||||
streamedCallsMu.Lock()
|
||||
advisorMessages = append([]chattest.OpenAIMessage(nil), req.Messages...)
|
||||
streamedCallsMu.Unlock()
|
||||
advisorCallSeen.Store(true)
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks(advisorReply)...,
|
||||
)
|
||||
default:
|
||||
// Parent turn 2: observe the advisor tool result and close
|
||||
// out with a final text reply.
|
||||
streamedCallsMu.Lock()
|
||||
finalCallMessages = append([]chattest.OpenAIMessage(nil), req.Messages...)
|
||||
streamedCallsMu.Unlock()
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("acknowledged")...,
|
||||
)
|
||||
}
|
||||
})
|
||||
|
||||
user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
|
||||
seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{
|
||||
Enabled: true,
|
||||
MaxUsesPerRun: 3,
|
||||
MaxOutputTokens: 16384,
|
||||
})
|
||||
server := newActiveTestServer(t, db, ps)
|
||||
|
||||
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
Title: "advisor-happy-path",
|
||||
ModelConfigID: model.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText("help me refactor this module"),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
if got.Status != database.ChatStatusWaiting &&
|
||||
got.Status != database.ChatStatusError {
|
||||
return false
|
||||
}
|
||||
return streamedCallCount.Load() >= 3
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
streamedCallsMu.Lock()
|
||||
gotAdvisorMessages := append([]chattest.OpenAIMessage(nil), advisorMessages...)
|
||||
gotFinalMessages := append([]chattest.OpenAIMessage(nil), finalCallMessages...)
|
||||
streamedCallsMu.Unlock()
|
||||
|
||||
require.True(t, advisorCallSeen.Load(),
|
||||
"the nested advisor call must execute; missing it means the tool never ran")
|
||||
require.NotEmpty(t, gotAdvisorMessages,
|
||||
"advisor call must receive the nested prompt messages")
|
||||
require.NotEmpty(t, gotFinalMessages,
|
||||
"parent must make a follow-up call after the advisor result")
|
||||
|
||||
var advisorSawQuestion bool
|
||||
var advisorSawUserTurn bool
|
||||
for _, msg := range gotAdvisorMessages {
|
||||
if strings.Contains(msg.Content, "how should I approach this refactor?") {
|
||||
advisorSawQuestion = true
|
||||
}
|
||||
if msg.Role == "user" && strings.Contains(msg.Content, "help me refactor this module") {
|
||||
advisorSawUserTurn = true
|
||||
}
|
||||
}
|
||||
require.True(t, advisorSawQuestion,
|
||||
"advisor must receive the parent's question verbatim")
|
||||
require.True(t, advisorSawUserTurn,
|
||||
"advisor must receive the parent's conversation snapshot as nested context")
|
||||
|
||||
for _, msg := range gotAdvisorMessages {
|
||||
require.NotContains(t, msg.Content, chatadvisor.ParentGuidanceBlock,
|
||||
"ParentGuidanceBlock must be stripped before reaching the advisor")
|
||||
}
|
||||
|
||||
var parentSawAdvisorResult bool
|
||||
for _, msg := range gotFinalMessages {
|
||||
if msg.Role == "tool" && strings.Contains(msg.Content, advisorReply) {
|
||||
parentSawAdvisorResult = true
|
||||
break
|
||||
}
|
||||
}
|
||||
require.True(t, parentSawAdvisorResult,
|
||||
"parent must see the advisor reply in its continuation call")
|
||||
}
|
||||
|
||||
// TestAdvisorGating_ChildChat guards the second dimension of the advisor
|
||||
// eligibility condition: even with advisor enabled, a chat whose
|
||||
// ParentChatID is set must not register the advisor tool or receive the
|
||||
// advisor guidance block. Without this coverage, a refactor that removes
|
||||
// or weakens the !chat.ParentChatID.Valid guard would leak advisor into
|
||||
// child chats, and the recursive advisor-inside-subagent cost risk the
|
||||
// guard exists to prevent would ship silently.
|
||||
//
|
||||
// The earlier version of this test drove the gating path through
|
||||
// spawn_agent, which made it dependent on subagent wiring that changed
|
||||
// repeatedly upstream. This version seeds the parent chat directly in the
|
||||
// database and asks the server to create a child chat with a valid
|
||||
// ParentChatID, exercising the same gating path with no subagent tooling
|
||||
// in the way.
|
||||
func TestAdvisorGating_ChildChat(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
var toolsMu sync.Mutex
|
||||
var capturedTools []string
|
||||
var capturedMessages []chattest.OpenAIMessage
|
||||
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
|
||||
names := make([]string, 0, len(req.Tools))
|
||||
for _, tool := range req.Tools {
|
||||
names = append(names, tool.Function.Name)
|
||||
}
|
||||
toolsMu.Lock()
|
||||
capturedTools = names
|
||||
capturedMessages = append([]chattest.OpenAIMessage(nil), req.Messages...)
|
||||
toolsMu.Unlock()
|
||||
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("done")...,
|
||||
)
|
||||
})
|
||||
|
||||
user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
|
||||
seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{
|
||||
Enabled: true,
|
||||
MaxUsesPerRun: 3,
|
||||
MaxOutputTokens: 16384,
|
||||
})
|
||||
|
||||
// Seed the parent chat directly in the database so the test server
|
||||
// never executes the root turn. That keeps this test focused on the
|
||||
// child-chat gating path without depending on subagent wiring.
|
||||
parent, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
Status: database.ChatStatusWaiting,
|
||||
ClientType: database.ChatClientTypeUi,
|
||||
LastModelConfigID: model.ID,
|
||||
Title: "advisor-root-parent",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
server := newActiveTestServer(t, db, ps)
|
||||
|
||||
childChat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
Title: "advisor-child",
|
||||
ModelConfigID: model.ID,
|
||||
ParentChatID: uuid.NullUUID{UUID: parent.ID, Valid: true},
|
||||
RootChatID: uuid.NullUUID{UUID: parent.ID, Valid: true},
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText("hi"),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := db.GetChatByID(ctx, childChat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
return got.Status == database.ChatStatusWaiting ||
|
||||
got.Status == database.ChatStatusError
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
toolsMu.Lock()
|
||||
tools := append([]string(nil), capturedTools...)
|
||||
messages := append([]chattest.OpenAIMessage(nil), capturedMessages...)
|
||||
toolsMu.Unlock()
|
||||
|
||||
require.NotEmpty(t, messages, "expected a streamed LLM request for the child chat")
|
||||
require.NotContains(t, tools, chatadvisor.ToolName,
|
||||
"advisor tool must not be registered for child chats even when enabled")
|
||||
for _, msg := range messages {
|
||||
require.NotContains(t, msg.Content, chatadvisor.ParentGuidanceBlock,
|
||||
"child chat must not contain advisor guidance")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdvisorGating_PlanMode guards the third dimension of the advisor
|
||||
// eligibility condition: plan-mode turns must not register the advisor tool
|
||||
// or inject the parent guidance block. Without this test, deleting the
|
||||
// !isPlanModeTurn guard would still leave the other two gating tests green
|
||||
// even though advisor would now leak into plan mode.
|
||||
func TestAdvisorGating_PlanMode(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
var toolsMu sync.Mutex
|
||||
var capturedTools []string
|
||||
var capturedMessages []chattest.OpenAIMessage
|
||||
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
|
||||
names := make([]string, 0, len(req.Tools))
|
||||
for _, tool := range req.Tools {
|
||||
names = append(names, tool.Function.Name)
|
||||
}
|
||||
toolsMu.Lock()
|
||||
capturedTools = names
|
||||
capturedMessages = append([]chattest.OpenAIMessage(nil), req.Messages...)
|
||||
toolsMu.Unlock()
|
||||
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("plan mode reply")...,
|
||||
)
|
||||
})
|
||||
|
||||
user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
|
||||
seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{
|
||||
Enabled: true,
|
||||
MaxUsesPerRun: 3,
|
||||
MaxOutputTokens: 16384,
|
||||
})
|
||||
server := newActiveTestServer(t, db, ps)
|
||||
|
||||
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
Title: "advisor-plan-mode",
|
||||
ModelConfigID: model.ID,
|
||||
PlanMode: database.NullChatPlanMode{ChatPlanMode: database.ChatPlanModePlan, Valid: true},
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText("draft a plan"),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
return got.Status == database.ChatStatusWaiting ||
|
||||
got.Status == database.ChatStatusError
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
toolsMu.Lock()
|
||||
tools := append([]string(nil), capturedTools...)
|
||||
messages := append([]chattest.OpenAIMessage(nil), capturedMessages...)
|
||||
toolsMu.Unlock()
|
||||
|
||||
require.NotEmpty(t, messages, "expected a streamed LLM request")
|
||||
require.NotContains(t, tools, "advisor",
|
||||
"plan-mode turns must not register the advisor tool even when enabled")
|
||||
for _, msg := range messages {
|
||||
require.NotContains(t, msg.Content, chatadvisor.ParentGuidanceBlock,
|
||||
"plan-mode turns must not inject advisor guidance")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdvisorGating_ExploreSubagent guards the fourth dimension of the
|
||||
// advisor eligibility condition: Explore chats (root or subagent) run
|
||||
// under allowedExploreToolNames, whose policy does not include advisor,
|
||||
// so the runtime must not register the advisor tool or inject the
|
||||
// parent guidance block there. Without this test, deleting the
|
||||
// !isExploreSubagent guard would leave the other gating tests green
|
||||
// while leaking advisor into explore chats.
|
||||
func TestAdvisorGating_ExploreSubagent(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
var toolsMu sync.Mutex
|
||||
var capturedTools []string
|
||||
var capturedMessages []chattest.OpenAIMessage
|
||||
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
|
||||
names := make([]string, 0, len(req.Tools))
|
||||
for _, tool := range req.Tools {
|
||||
names = append(names, tool.Function.Name)
|
||||
}
|
||||
toolsMu.Lock()
|
||||
capturedTools = names
|
||||
capturedMessages = append([]chattest.OpenAIMessage(nil), req.Messages...)
|
||||
toolsMu.Unlock()
|
||||
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("explore reply")...,
|
||||
)
|
||||
})
|
||||
|
||||
user, org, model := seedChatDependenciesWithProvider(ctx, t, db, "openai-compat", openAIURL)
|
||||
seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{
|
||||
Enabled: true,
|
||||
MaxUsesPerRun: 3,
|
||||
MaxOutputTokens: 16384,
|
||||
})
|
||||
server := newActiveTestServer(t, db, ps)
|
||||
|
||||
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
Title: "advisor-explore",
|
||||
ModelConfigID: model.ID,
|
||||
ChatMode: database.NullChatMode{
|
||||
ChatMode: database.ChatModeExplore,
|
||||
Valid: true,
|
||||
},
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText("inspect the codebase"),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
got, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
return got.Status == database.ChatStatusWaiting ||
|
||||
got.Status == database.ChatStatusError
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
toolsMu.Lock()
|
||||
tools := append([]string(nil), capturedTools...)
|
||||
messages := append([]chattest.OpenAIMessage(nil), capturedMessages...)
|
||||
toolsMu.Unlock()
|
||||
|
||||
require.NotEmpty(t, messages, "expected a streamed LLM request")
|
||||
require.NotContains(t, tools, chatadvisor.ToolName,
|
||||
"explore chats must not register the advisor tool even when enabled")
|
||||
for _, msg := range messages {
|
||||
require.NotContains(t, msg.Content, chatadvisor.ParentGuidanceBlock,
|
||||
"explore chats must not inject advisor guidance")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdvisorChainMode_SnapshotKeepsFullHistory exercises the advisor
|
||||
// runtime together with chain mode and asserts the snapshot captured for
|
||||
// the nested advisor call retains the full pre-chain prompt. Chain mode
|
||||
// otherwise strips assistant and tool turns from the prompt the outer
|
||||
// loop sees, so a regression that moves setAdvisorPromptSnapshot behind
|
||||
// filterPromptForChainMode, or drops the !chainModeActive guards in
|
||||
// PrepareMessages, would leak the filtered view into the advisor's
|
||||
// nested call. The advisor would then only see the trailing user
|
||||
// message, losing the context the outer model had been building on.
|
||||
func TestAdvisorChainMode_SnapshotKeepsFullHistory(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, ps := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitLong)
|
||||
|
||||
const (
|
||||
turn1User = "help me refactor this module"
|
||||
turn1Reply = "happy to help, tell me more"
|
||||
turn1RespID = "resp_turn1_advisor_chain"
|
||||
turn2User = "follow up question"
|
||||
advisorReply = "narrow the scope to one module"
|
||||
finalReply = "acknowledged"
|
||||
)
|
||||
|
||||
var (
|
||||
requestsMu sync.Mutex
|
||||
requests []recordedOpenAIRequest
|
||||
advisorRequestRaw []byte
|
||||
advisorCallSeen atomic.Bool
|
||||
)
|
||||
|
||||
openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
if !req.Stream {
|
||||
return chattest.OpenAINonStreamingResponse("title")
|
||||
}
|
||||
|
||||
// The advisor's nested call runs with no tools (MaxSteps=1,
|
||||
// empty tool set). Parent calls always carry the chat's tool
|
||||
// set, which includes the advisor tool.
|
||||
isAdvisorNested := len(req.Tools) == 0
|
||||
|
||||
requestsMu.Lock()
|
||||
requests = append(requests, recordOpenAIRequest(req))
|
||||
if isAdvisorNested {
|
||||
advisorRequestRaw = append([]byte(nil), req.RawBody...)
|
||||
advisorCallSeen.Store(true)
|
||||
}
|
||||
requestsMu.Unlock()
|
||||
|
||||
if isAdvisorNested {
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks(advisorReply)...,
|
||||
)
|
||||
}
|
||||
|
||||
// Turn 1 parent request: no previous_response_id yet, so chain
|
||||
// mode cannot activate. Respond with a plain text reply and
|
||||
// tag the stored response id so turn 2 can chain off it.
|
||||
if req.PreviousResponseID == nil {
|
||||
resp := chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks(turn1Reply)...,
|
||||
)
|
||||
resp.ResponseID = turn1RespID
|
||||
return resp
|
||||
}
|
||||
|
||||
// Turn 2 parent: chain mode is active. On the first pass call
|
||||
// advisor; on the continuation after the tool result arrives,
|
||||
// close out with a final text reply.
|
||||
var hasAdvisorResult bool
|
||||
for _, m := range req.Messages {
|
||||
if m.Role == "tool" && strings.Contains(m.Content, advisorReply) {
|
||||
hasAdvisorResult = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasAdvisorResult {
|
||||
return chattest.OpenAIStreamingResponse(chattest.OpenAIToolCallChunk(
|
||||
"advisor",
|
||||
`{"question":"should I keep going?"}`,
|
||||
))
|
||||
}
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks(finalReply)...,
|
||||
)
|
||||
})
|
||||
|
||||
user, org, _ := seedChatDependenciesWithProvider(ctx, t, db, "openai", openAIURL)
|
||||
storeEnabled := true
|
||||
// The OpenAI Responses API is the only provider code path where
|
||||
// chain mode activates. Store=true is the switch that routes this
|
||||
// provider/model through the Responses API and lets
|
||||
// IsResponsesStoreEnabled return true.
|
||||
responsesModel := insertChatModelConfigWithCallConfig(
|
||||
ctx, t, db, user.ID, "openai", "gpt-4o",
|
||||
codersdk.ChatModelCallConfig{
|
||||
ProviderOptions: &codersdk.ChatModelProviderOptions{
|
||||
OpenAI: &codersdk.ChatModelOpenAIProviderOptions{
|
||||
Store: &storeEnabled,
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
seedAdvisorConfig(ctx, t, db, codersdk.AdvisorConfig{
|
||||
Enabled: true,
|
||||
MaxUsesPerRun: 3,
|
||||
MaxOutputTokens: 16384,
|
||||
})
|
||||
server := newActiveTestServer(t, db, ps)
|
||||
|
||||
chat, err := server.CreateChat(ctx, chatd.CreateOptions{
|
||||
OrganizationID: org.ID,
|
||||
OwnerID: user.ID,
|
||||
Title: "advisor-chain-mode",
|
||||
ModelConfigID: responsesModel.ID,
|
||||
InitialUserContent: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText(turn1User),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
// Turn 1 must settle before turn 2 starts so the assistant row
|
||||
// with ProviderResponseID is visible to resolveChainMode.
|
||||
waitForChatProcessed(ctx, t, db, chat.ID, server)
|
||||
turn1Chat, err := db.GetChatByID(ctx, chat.ID)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, database.ChatStatusWaiting, turn1Chat.Status,
|
||||
"turn 1 must complete before turn 2 can be sent; last_error=%q", turn1Chat.LastError.String)
|
||||
|
||||
_, err = server.SendMessage(ctx, chatd.SendMessageOptions{
|
||||
ChatID: chat.ID,
|
||||
CreatedBy: user.ID,
|
||||
Content: []codersdk.ChatMessagePart{
|
||||
codersdk.ChatMessageText(turn2User),
|
||||
},
|
||||
})
|
||||
require.NoError(t, err)
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
if !advisorCallSeen.Load() {
|
||||
return false
|
||||
}
|
||||
got, getErr := db.GetChatByID(ctx, chat.ID)
|
||||
if getErr != nil {
|
||||
return false
|
||||
}
|
||||
return got.Status == database.ChatStatusWaiting ||
|
||||
got.Status == database.ChatStatusError
|
||||
}, testutil.WaitLong, testutil.IntervalFast)
|
||||
|
||||
requestsMu.Lock()
|
||||
gotAdvisorBody := append([]byte(nil), advisorRequestRaw...)
|
||||
gotRequests := append([]recordedOpenAIRequest(nil), requests...)
|
||||
requestsMu.Unlock()
|
||||
|
||||
// Chain mode must have actually fired on turn 2, otherwise this
|
||||
// test degenerates to TestAdvisorHappyPath_RootChat.
|
||||
var chainModeActivated bool
|
||||
for _, r := range gotRequests {
|
||||
if r.PreviousResponseID != nil && *r.PreviousResponseID == turn1RespID {
|
||||
chainModeActivated = true
|
||||
break
|
||||
}
|
||||
}
|
||||
require.True(t, chainModeActivated,
|
||||
"turn 2 parent request must carry previous_response_id; without it this test does not exercise chain mode")
|
||||
|
||||
require.True(t, advisorCallSeen.Load(),
|
||||
"the nested advisor call must execute under chain mode")
|
||||
require.NotEmpty(t, gotAdvisorBody,
|
||||
"advisor call must receive a non-empty request body")
|
||||
|
||||
// The core assertion: the advisor snapshot must retain turn 1
|
||||
// context. Chain mode filtering strips assistant and tool turns
|
||||
// from the prompt the outer loop sees, so if that filtered view
|
||||
// leaked into the snapshot the advisor would only see turn 2's
|
||||
// trailing user message. The advisor's nested call goes through
|
||||
// the OpenAI Responses API, which encodes its prompt in the
|
||||
// "input" field rather than "messages", so we inspect the raw
|
||||
// request body for both turn-1 substrings.
|
||||
require.Contains(t, string(gotAdvisorBody), turn1User,
|
||||
"advisor snapshot must retain the turn 1 user message even when chain mode is active")
|
||||
require.Contains(t, string(gotAdvisorBody), turn1Reply,
|
||||
"advisor snapshot must retain the turn 1 assistant message even when chain mode is active")
|
||||
}
|
||||
|
||||
func seedAdvisorConfig(
|
||||
ctx context.Context,
|
||||
t *testing.T,
|
||||
db database.Store,
|
||||
cfg codersdk.AdvisorConfig,
|
||||
) {
|
||||
t.Helper()
|
||||
|
||||
data, err := json.Marshal(cfg)
|
||||
require.NoError(t, err)
|
||||
err = db.UpsertChatAdvisorConfig(
|
||||
dbauthz.AsSystemRestricted(ctx),
|
||||
string(data),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
@@ -720,6 +720,145 @@ func ReasoningEffortFromChat(provider string, value *string) *string {
|
||||
}
|
||||
}
|
||||
|
||||
// ApplyReasoningEffortToOptions applies the given reasoning_effort to every
|
||||
// provider entry in providerOptions that understands it. When model is
|
||||
// non-nil and the options map has no entry for the model's provider, this
|
||||
// function seeds a minimal provider-specific options struct so the mutation
|
||||
// still lands. Callers that produced providerOptions from a chat model
|
||||
// config with no provider_options block would otherwise see
|
||||
// reasoning_effort silently dropped.
|
||||
//
|
||||
// The returned map is the (possibly newly-allocated) providerOptions; the
|
||||
// input is mutated in-place when non-nil.
|
||||
func ApplyReasoningEffortToOptions(
|
||||
providerOptions fantasy.ProviderOptions,
|
||||
model fantasy.LanguageModel,
|
||||
reasoningEffort string,
|
||||
) fantasy.ProviderOptions {
|
||||
reasoningEffort = strings.TrimSpace(reasoningEffort)
|
||||
if reasoningEffort == "" {
|
||||
return providerOptions
|
||||
}
|
||||
|
||||
if model != nil {
|
||||
providerOptions = seedProviderOptionsForModel(providerOptions, model)
|
||||
}
|
||||
if providerOptions == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
applyReasoningEffortDispatch(providerOptions, reasoningEffort)
|
||||
return providerOptions
|
||||
}
|
||||
|
||||
// seedProviderOptionsForModel ensures providerOptions has an entry for the
|
||||
// given model's provider, allocating a minimal options struct when absent.
|
||||
// Returns the possibly newly-allocated options map. Unknown providers are
|
||||
// left untouched so callers get their input back unchanged.
|
||||
func seedProviderOptionsForModel(
|
||||
providerOptions fantasy.ProviderOptions,
|
||||
model fantasy.LanguageModel,
|
||||
) fantasy.ProviderOptions {
|
||||
provider := model.Provider()
|
||||
var seed fantasy.ProviderOptionsData
|
||||
switch provider {
|
||||
case fantasyopenai.Name:
|
||||
if fantasyopenai.IsResponsesModel(model.Model()) {
|
||||
seed = &fantasyopenai.ResponsesProviderOptions{}
|
||||
} else {
|
||||
seed = &fantasyopenai.ProviderOptions{}
|
||||
}
|
||||
case fantasyanthropic.Name:
|
||||
seed = &fantasyanthropic.ProviderOptions{}
|
||||
case fantasyopenaicompat.Name:
|
||||
seed = &fantasyopenaicompat.ProviderOptions{}
|
||||
case fantasyopenrouter.Name:
|
||||
seed = &fantasyopenrouter.ProviderOptions{}
|
||||
case fantasyvercel.Name:
|
||||
seed = &fantasyvercel.ProviderOptions{}
|
||||
default:
|
||||
return providerOptions
|
||||
}
|
||||
|
||||
if providerOptions == nil {
|
||||
providerOptions = fantasy.ProviderOptions{}
|
||||
}
|
||||
if _, ok := providerOptions[provider]; !ok {
|
||||
providerOptions[provider] = seed
|
||||
}
|
||||
return providerOptions
|
||||
}
|
||||
|
||||
// applyReasoningEffortDispatch routes the normalized reasoning_effort to
|
||||
// every provider entry present in providerOptions. Adding a new provider
|
||||
// here (and only here) keeps chatd callers in sync automatically.
|
||||
func applyReasoningEffortDispatch(
|
||||
providerOptions fantasy.ProviderOptions,
|
||||
reasoningEffort string,
|
||||
) {
|
||||
if normalized := ReasoningEffortFromChat(
|
||||
fantasyopenai.Name,
|
||||
&reasoningEffort,
|
||||
); normalized != nil {
|
||||
effort := fantasyopenai.ReasoningEffort(*normalized)
|
||||
if raw, ok := providerOptions[fantasyopenai.Name]; ok {
|
||||
switch opts := raw.(type) {
|
||||
case *fantasyopenai.ProviderOptions:
|
||||
opts.ReasoningEffort = &effort
|
||||
case *fantasyopenai.ResponsesProviderOptions:
|
||||
opts.ReasoningEffort = &effort
|
||||
}
|
||||
}
|
||||
if raw, ok := providerOptions[fantasyopenaicompat.Name]; ok {
|
||||
if opts, ok := raw.(*fantasyopenaicompat.ProviderOptions); ok {
|
||||
opts.ReasoningEffort = &effort
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if normalized := ReasoningEffortFromChat(
|
||||
fantasyanthropic.Name,
|
||||
&reasoningEffort,
|
||||
); normalized != nil {
|
||||
if raw, ok := providerOptions[fantasyanthropic.Name]; ok {
|
||||
if opts, ok := raw.(*fantasyanthropic.ProviderOptions); ok {
|
||||
effort := fantasyanthropic.Effort(*normalized)
|
||||
opts.Effort = &effort
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if normalized := ReasoningEffortFromChat(
|
||||
fantasyopenrouter.Name,
|
||||
&reasoningEffort,
|
||||
); normalized != nil {
|
||||
if raw, ok := providerOptions[fantasyopenrouter.Name]; ok {
|
||||
if opts, ok := raw.(*fantasyopenrouter.ProviderOptions); ok {
|
||||
if opts.Reasoning == nil {
|
||||
opts.Reasoning = &fantasyopenrouter.ReasoningOptions{}
|
||||
}
|
||||
effort := fantasyopenrouter.ReasoningEffort(*normalized)
|
||||
opts.Reasoning.Effort = &effort
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if normalized := ReasoningEffortFromChat(
|
||||
fantasyvercel.Name,
|
||||
&reasoningEffort,
|
||||
); normalized != nil {
|
||||
if raw, ok := providerOptions[fantasyvercel.Name]; ok {
|
||||
if opts, ok := raw.(*fantasyvercel.ProviderOptions); ok {
|
||||
if opts.Reasoning == nil {
|
||||
opts.Reasoning = &fantasyvercel.ReasoningOptions{}
|
||||
}
|
||||
effort := fantasyvercel.ReasoningEffort(*normalized)
|
||||
opts.Reasoning.Effort = &effort
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// OpenAITextVerbosityFromChat normalizes chat-config text verbosity values for
|
||||
// OpenAI and returns the canonical provider verbosity value.
|
||||
func OpenAITextVerbosityFromChat(value *string) *fantasyopenai.TextVerbosity {
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
fantasyanthropic "charm.land/fantasy/providers/anthropic"
|
||||
fantasybedrock "charm.land/fantasy/providers/bedrock"
|
||||
fantasyopenai "charm.land/fantasy/providers/openai"
|
||||
fantasyopenaicompat "charm.land/fantasy/providers/openaicompat"
|
||||
fantasyopenrouter "charm.land/fantasy/providers/openrouter"
|
||||
fantasyvercel "charm.land/fantasy/providers/vercel"
|
||||
"github.com/google/uuid"
|
||||
@@ -1391,3 +1392,212 @@ func TestMergeMissingProviderOptions_OpenRouterNested(t *testing.T) {
|
||||
require.Equal(t, []string{"int8"}, options.OpenRouter.Provider.Quantizations)
|
||||
require.Equal(t, "latency", *options.OpenRouter.Provider.Sort)
|
||||
}
|
||||
|
||||
// TestApplyReasoningEffortToOptions covers every provider's mutation branch
|
||||
// plus the seeding path for missing provider entries. A typo or wrong type
|
||||
// assertion in any branch fails a unit test here rather than silently
|
||||
// dropping the admin-configured reasoning effort in chatd callers.
|
||||
func TestApplyReasoningEffortToOptions(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
t.Run("NilOptionsAndNilModelIsNoOp", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
// Must not panic when options and model are both nil.
|
||||
got := chatprovider.ApplyReasoningEffortToOptions(nil, nil, "medium")
|
||||
require.Nil(t, got)
|
||||
})
|
||||
|
||||
t.Run("EmptyEffortReturnsInputUnchanged", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
model := &chattest.FakeModel{ProviderName: fantasyopenai.Name, ModelName: "gpt-4"}
|
||||
got := chatprovider.ApplyReasoningEffortToOptions(nil, model, " ")
|
||||
require.Nil(t, got)
|
||||
})
|
||||
|
||||
t.Run("EmptyEffortPreservesExistingOptions", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
effort := fantasyopenai.ReasoningEffortLow
|
||||
opts := &fantasyopenai.ProviderOptions{ReasoningEffort: &effort}
|
||||
providerOptions := fantasy.ProviderOptions{fantasyopenai.Name: opts}
|
||||
|
||||
got := chatprovider.ApplyReasoningEffortToOptions(providerOptions, nil, "")
|
||||
require.NotNil(t, opts.ReasoningEffort)
|
||||
require.Equal(t, fantasyopenai.ReasoningEffortLow, *opts.ReasoningEffort)
|
||||
// The input map must be returned untouched rather than allocated anew.
|
||||
require.Len(t, got, 1)
|
||||
})
|
||||
|
||||
t.Run("UnrecognizedEffortLeavesOptionsUntouched", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
opts := &fantasyopenai.ProviderOptions{}
|
||||
providerOptions := fantasy.ProviderOptions{fantasyopenai.Name: opts}
|
||||
|
||||
chatprovider.ApplyReasoningEffortToOptions(providerOptions, nil, "not-a-real-effort")
|
||||
require.Nil(t, opts.ReasoningEffort)
|
||||
})
|
||||
|
||||
t.Run("OpenAIProviderOptions", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
opts := &fantasyopenai.ProviderOptions{}
|
||||
providerOptions := fantasy.ProviderOptions{fantasyopenai.Name: opts}
|
||||
|
||||
chatprovider.ApplyReasoningEffortToOptions(providerOptions, nil, "medium")
|
||||
require.NotNil(t, opts.ReasoningEffort)
|
||||
require.Equal(t, fantasyopenai.ReasoningEffortMedium, *opts.ReasoningEffort)
|
||||
})
|
||||
|
||||
t.Run("OpenAIResponsesProviderOptions", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
opts := &fantasyopenai.ResponsesProviderOptions{}
|
||||
providerOptions := fantasy.ProviderOptions{fantasyopenai.Name: opts}
|
||||
|
||||
chatprovider.ApplyReasoningEffortToOptions(providerOptions, nil, "medium")
|
||||
require.NotNil(t, opts.ReasoningEffort)
|
||||
require.Equal(t, fantasyopenai.ReasoningEffortMedium, *opts.ReasoningEffort)
|
||||
})
|
||||
|
||||
t.Run("OpenAICompatProviderOptions", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
opts := &fantasyopenaicompat.ProviderOptions{}
|
||||
providerOptions := fantasy.ProviderOptions{fantasyopenaicompat.Name: opts}
|
||||
|
||||
chatprovider.ApplyReasoningEffortToOptions(providerOptions, nil, "medium")
|
||||
require.NotNil(t, opts.ReasoningEffort)
|
||||
require.Equal(t, fantasyopenai.ReasoningEffortMedium, *opts.ReasoningEffort)
|
||||
})
|
||||
|
||||
t.Run("AnthropicProviderOptions", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
opts := &fantasyanthropic.ProviderOptions{}
|
||||
providerOptions := fantasy.ProviderOptions{fantasyanthropic.Name: opts}
|
||||
|
||||
chatprovider.ApplyReasoningEffortToOptions(providerOptions, nil, "high")
|
||||
require.NotNil(t, opts.Effort)
|
||||
require.Equal(t, fantasyanthropic.EffortHigh, *opts.Effort)
|
||||
})
|
||||
|
||||
t.Run("OpenRouterAllocatesReasoningOptions", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
opts := &fantasyopenrouter.ProviderOptions{}
|
||||
providerOptions := fantasy.ProviderOptions{fantasyopenrouter.Name: opts}
|
||||
|
||||
chatprovider.ApplyReasoningEffortToOptions(providerOptions, nil, "medium")
|
||||
require.NotNil(t, opts.Reasoning, "Reasoning container must be allocated")
|
||||
require.NotNil(t, opts.Reasoning.Effort)
|
||||
require.Equal(t, fantasyopenrouter.ReasoningEffort("medium"), *opts.Reasoning.Effort)
|
||||
})
|
||||
|
||||
t.Run("OpenRouterPreservesExistingReasoningContainer", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
enabled := true
|
||||
opts := &fantasyopenrouter.ProviderOptions{
|
||||
Reasoning: &fantasyopenrouter.ReasoningOptions{Enabled: &enabled},
|
||||
}
|
||||
providerOptions := fantasy.ProviderOptions{fantasyopenrouter.Name: opts}
|
||||
|
||||
chatprovider.ApplyReasoningEffortToOptions(providerOptions, nil, "high")
|
||||
require.NotNil(t, opts.Reasoning.Enabled)
|
||||
require.True(t, *opts.Reasoning.Enabled)
|
||||
require.NotNil(t, opts.Reasoning.Effort)
|
||||
require.Equal(t, fantasyopenrouter.ReasoningEffort("high"), *opts.Reasoning.Effort)
|
||||
})
|
||||
|
||||
t.Run("VercelAllocatesReasoningOptions", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
opts := &fantasyvercel.ProviderOptions{}
|
||||
providerOptions := fantasy.ProviderOptions{fantasyvercel.Name: opts}
|
||||
|
||||
chatprovider.ApplyReasoningEffortToOptions(providerOptions, nil, "minimal")
|
||||
require.NotNil(t, opts.Reasoning)
|
||||
require.NotNil(t, opts.Reasoning.Effort)
|
||||
require.Equal(t, fantasyvercel.ReasoningEffortMinimal, *opts.Reasoning.Effort)
|
||||
})
|
||||
|
||||
t.Run("MultipleProvidersReceiveMutations", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
openaiOpts := &fantasyopenai.ProviderOptions{}
|
||||
anthropicOpts := &fantasyanthropic.ProviderOptions{}
|
||||
providerOptions := fantasy.ProviderOptions{
|
||||
fantasyopenai.Name: openaiOpts,
|
||||
fantasyanthropic.Name: anthropicOpts,
|
||||
}
|
||||
|
||||
chatprovider.ApplyReasoningEffortToOptions(providerOptions, nil, "high")
|
||||
require.NotNil(t, openaiOpts.ReasoningEffort)
|
||||
require.Equal(t, fantasyopenai.ReasoningEffortHigh, *openaiOpts.ReasoningEffort)
|
||||
require.NotNil(t, anthropicOpts.Effort)
|
||||
require.Equal(t, fantasyanthropic.EffortHigh, *anthropicOpts.Effort)
|
||||
})
|
||||
|
||||
t.Run("SeedsOpenAICompletionsWhenModelHasNoOptions", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
// A model name absent from the Responses allowlist must seed
|
||||
// the completions options struct so reasoning_effort lands.
|
||||
model := &chattest.FakeModel{ProviderName: fantasyopenai.Name, ModelName: "not-a-real-openai-model"}
|
||||
got := chatprovider.ApplyReasoningEffortToOptions(nil, model, "medium")
|
||||
require.NotNil(t, got)
|
||||
opts, ok := got[fantasyopenai.Name].(*fantasyopenai.ProviderOptions)
|
||||
require.True(t, ok, "expected *ProviderOptions for non-Responses model, got %T", got[fantasyopenai.Name])
|
||||
require.NotNil(t, opts.ReasoningEffort)
|
||||
require.Equal(t, fantasyopenai.ReasoningEffortMedium, *opts.ReasoningEffort)
|
||||
})
|
||||
|
||||
t.Run("SeedsOpenAIResponsesWhenModelIsResponsesModel", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
// A model name in the Responses allowlist must seed the
|
||||
// Responses-specific options struct so the provider routes to
|
||||
// the Responses endpoint.
|
||||
model := &chattest.FakeModel{ProviderName: fantasyopenai.Name, ModelName: "gpt-4"}
|
||||
got := chatprovider.ApplyReasoningEffortToOptions(nil, model, "medium")
|
||||
require.NotNil(t, got)
|
||||
opts, ok := got[fantasyopenai.Name].(*fantasyopenai.ResponsesProviderOptions)
|
||||
require.True(t, ok, "expected *ResponsesProviderOptions for Responses model, got %T", got[fantasyopenai.Name])
|
||||
require.NotNil(t, opts.ReasoningEffort)
|
||||
require.Equal(t, fantasyopenai.ReasoningEffortMedium, *opts.ReasoningEffort)
|
||||
})
|
||||
|
||||
t.Run("SeedsAnthropicWhenModelHasNoOptions", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
model := &chattest.FakeModel{ProviderName: fantasyanthropic.Name, ModelName: "claude-3-5"}
|
||||
got := chatprovider.ApplyReasoningEffortToOptions(nil, model, "high")
|
||||
require.NotNil(t, got)
|
||||
opts, ok := got[fantasyanthropic.Name].(*fantasyanthropic.ProviderOptions)
|
||||
require.True(t, ok)
|
||||
require.NotNil(t, opts.Effort)
|
||||
require.Equal(t, fantasyanthropic.EffortHigh, *opts.Effort)
|
||||
})
|
||||
|
||||
t.Run("SeedsOpenRouterWhenModelHasNoOptions", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
model := &chattest.FakeModel{ProviderName: fantasyopenrouter.Name, ModelName: "openrouter-x"}
|
||||
got := chatprovider.ApplyReasoningEffortToOptions(nil, model, "low")
|
||||
require.NotNil(t, got)
|
||||
opts, ok := got[fantasyopenrouter.Name].(*fantasyopenrouter.ProviderOptions)
|
||||
require.True(t, ok)
|
||||
require.NotNil(t, opts.Reasoning)
|
||||
require.NotNil(t, opts.Reasoning.Effort)
|
||||
require.Equal(t, fantasyopenrouter.ReasoningEffort("low"), *opts.Reasoning.Effort)
|
||||
})
|
||||
|
||||
t.Run("UnknownProviderReturnsInputUnchanged", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
model := &chattest.FakeModel{ProviderName: "unknown", ModelName: "x"}
|
||||
got := chatprovider.ApplyReasoningEffortToOptions(nil, model, "medium")
|
||||
require.Nil(t, got)
|
||||
})
|
||||
|
||||
t.Run("PreservesExistingProviderEntry", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
existing := &fantasyopenai.ProviderOptions{}
|
||||
existingEffort := fantasyopenai.ReasoningEffortLow
|
||||
existing.ReasoningEffort = &existingEffort
|
||||
providerOptions := fantasy.ProviderOptions{fantasyopenai.Name: existing}
|
||||
|
||||
model := &chattest.FakeModel{ProviderName: fantasyopenai.Name, ModelName: "gpt-4"}
|
||||
got := chatprovider.ApplyReasoningEffortToOptions(providerOptions, model, "medium")
|
||||
require.Same(t, existing, got[fantasyopenai.Name],
|
||||
"existing provider entry must not be replaced")
|
||||
// The reasoning effort on the existing entry is overwritten.
|
||||
require.Equal(t, fantasyopenai.ReasoningEffortMedium, *existing.ReasoningEffort)
|
||||
})
|
||||
}
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
package chattest
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
@@ -53,6 +55,10 @@ type OpenAIRequest struct {
|
||||
Prompt []interface{} `json:"prompt,omitempty"` // Responses API input or prompt.
|
||||
Store *bool `json:"store,omitempty"`
|
||||
PreviousResponseID *string `json:"previous_response_id,omitempty"`
|
||||
// RawBody holds the original request body so callers can inspect
|
||||
// fields the typed struct does not expose, such as the Responses
|
||||
// API "input" payload. It is populated before JSON decoding.
|
||||
RawBody []byte `json:"-"`
|
||||
// TODO: encoding/json ignores inline tags. Add custom UnmarshalJSON to capture unknown keys.
|
||||
Options map[string]interface{} `json:",inline"` //nolint:revive
|
||||
}
|
||||
@@ -205,12 +211,18 @@ func NewOpenAI(t testing.TB, handler OpenAIHandler) string {
|
||||
}
|
||||
|
||||
func (s *openAIServer) handleChatCompletions(w http.ResponseWriter, r *http.Request) {
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
var req OpenAIRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
if err := json.NewDecoder(bytes.NewReader(body)).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
req.Request = r
|
||||
req.RawBody = body
|
||||
|
||||
s.mu.Lock()
|
||||
s.request = &req
|
||||
@@ -221,12 +233,18 @@ func (s *openAIServer) handleChatCompletions(w http.ResponseWriter, r *http.Requ
|
||||
}
|
||||
|
||||
func (s *openAIServer) handleResponses(w http.ResponseWriter, r *http.Request) {
|
||||
body, err := io.ReadAll(r.Body)
|
||||
if err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
var req OpenAIRequest
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
if err := json.NewDecoder(bytes.NewReader(body)).Decode(&req); err != nil {
|
||||
http.Error(w, err.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
req.Request = r
|
||||
req.RawBody = body
|
||||
|
||||
s.mu.Lock()
|
||||
s.request = &req
|
||||
@@ -407,8 +425,18 @@ func writeResponsesAPIStreaming(t testing.TB, w http.ResponseWriter, r *http.Req
|
||||
responseModel := "gpt-4"
|
||||
sequenceNumber := int64(0)
|
||||
textOffset := 0
|
||||
itemIDs := make(map[int]string)
|
||||
itemTexts := make(map[int]string)
|
||||
// outputs tracks per-output-index state so the done-event emission
|
||||
// at stream close can distinguish message items (text) from
|
||||
// function_call items (tool invocation).
|
||||
type outputItemState struct {
|
||||
itemType string // "message" or "function_call"
|
||||
itemID string
|
||||
text string // accumulated text for message items
|
||||
callID string // call_id for function_call items
|
||||
toolName string // function name for function_call items
|
||||
arguments string // accumulated arguments for function_call items
|
||||
}
|
||||
outputs := make(map[int]*outputItemState)
|
||||
|
||||
writeEvent := func(eventType string, payload map[string]interface{}) bool {
|
||||
payload["type"] = eventType
|
||||
@@ -574,50 +602,73 @@ func writeResponsesAPIStreaming(t testing.TB, w http.ResponseWriter, r *http.Req
|
||||
return
|
||||
case chunk, ok = <-resp.StreamingChunks:
|
||||
if !ok {
|
||||
indices := make([]int, 0, len(itemIDs))
|
||||
for outputIndex := range itemIDs {
|
||||
indices := make([]int, 0, len(outputs))
|
||||
for outputIndex := range outputs {
|
||||
indices = append(indices, outputIndex)
|
||||
}
|
||||
sort.Ints(indices)
|
||||
for _, outputIndex := range indices {
|
||||
itemID := itemIDs[outputIndex]
|
||||
text := itemTexts[outputIndex]
|
||||
if !writeEvent("response.output_text.done", map[string]interface{}{
|
||||
"item_id": itemID,
|
||||
"output_index": outputIndex,
|
||||
"content_index": 0,
|
||||
"text": text,
|
||||
"logprobs": []interface{}{},
|
||||
}) {
|
||||
return
|
||||
}
|
||||
if !writeEvent("response.content_part.done", map[string]interface{}{
|
||||
"item_id": itemID,
|
||||
"output_index": outputIndex,
|
||||
"content_index": 0,
|
||||
"part": map[string]interface{}{
|
||||
"type": "output_text",
|
||||
"text": text,
|
||||
},
|
||||
}) {
|
||||
return
|
||||
}
|
||||
if !writeEvent("response.output_item.done", map[string]interface{}{
|
||||
"output_index": outputIndex,
|
||||
"item": map[string]interface{}{
|
||||
"type": "message",
|
||||
"id": itemID,
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": []interface{}{
|
||||
map[string]interface{}{
|
||||
"type": "output_text",
|
||||
"text": text,
|
||||
state := outputs[outputIndex]
|
||||
switch state.itemType {
|
||||
case "function_call":
|
||||
if !writeEvent("response.function_call_arguments.done", map[string]interface{}{
|
||||
"item_id": state.itemID,
|
||||
"output_index": outputIndex,
|
||||
"arguments": state.arguments,
|
||||
}) {
|
||||
return
|
||||
}
|
||||
if !writeEvent("response.output_item.done", map[string]interface{}{
|
||||
"output_index": outputIndex,
|
||||
"item": map[string]interface{}{
|
||||
"type": "function_call",
|
||||
"id": state.itemID,
|
||||
"status": "completed",
|
||||
"call_id": state.callID,
|
||||
"name": state.toolName,
|
||||
"arguments": state.arguments,
|
||||
},
|
||||
}) {
|
||||
return
|
||||
}
|
||||
default:
|
||||
if !writeEvent("response.output_text.done", map[string]interface{}{
|
||||
"item_id": state.itemID,
|
||||
"output_index": outputIndex,
|
||||
"content_index": 0,
|
||||
"text": state.text,
|
||||
"logprobs": []interface{}{},
|
||||
}) {
|
||||
return
|
||||
}
|
||||
if !writeEvent("response.content_part.done", map[string]interface{}{
|
||||
"item_id": state.itemID,
|
||||
"output_index": outputIndex,
|
||||
"content_index": 0,
|
||||
"part": map[string]interface{}{
|
||||
"type": "output_text",
|
||||
"text": state.text,
|
||||
},
|
||||
}) {
|
||||
return
|
||||
}
|
||||
if !writeEvent("response.output_item.done", map[string]interface{}{
|
||||
"output_index": outputIndex,
|
||||
"item": map[string]interface{}{
|
||||
"type": "message",
|
||||
"id": state.itemID,
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": []interface{}{
|
||||
map[string]interface{}{
|
||||
"type": "output_text",
|
||||
"text": state.text,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}) {
|
||||
return
|
||||
}) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
if !writeEvent("response.completed", map[string]interface{}{
|
||||
@@ -645,15 +696,64 @@ func writeResponsesAPIStreaming(t testing.TB, w http.ResponseWriter, r *http.Req
|
||||
outputIndex = choice.Index
|
||||
}
|
||||
outputIndex += textOffset
|
||||
itemID, found := itemIDs[outputIndex]
|
||||
|
||||
if len(choice.ToolCalls) > 0 {
|
||||
for _, tc := range choice.ToolCalls {
|
||||
// Each tool call within a chunk owns a distinct
|
||||
// output item, so discriminate by the streaming
|
||||
// tc.Index. Without this, multiple tool calls in
|
||||
// one chunk collide on outputIndex and later
|
||||
// calls inherit the first call's id and name.
|
||||
toolOutputIndex := outputIndex + tc.Index
|
||||
state, found := outputs[toolOutputIndex]
|
||||
if !found {
|
||||
state = &outputItemState{
|
||||
itemType: "function_call",
|
||||
itemID: fmt.Sprintf("fc_%s", uuid.New().String()[:8]),
|
||||
callID: tc.ID,
|
||||
toolName: tc.Function.Name,
|
||||
}
|
||||
outputs[toolOutputIndex] = state
|
||||
if !writeEvent("response.output_item.added", map[string]interface{}{
|
||||
"output_index": toolOutputIndex,
|
||||
"item": map[string]interface{}{
|
||||
"type": "function_call",
|
||||
"id": state.itemID,
|
||||
"status": "in_progress",
|
||||
"call_id": state.callID,
|
||||
"name": state.toolName,
|
||||
"arguments": "",
|
||||
},
|
||||
}) {
|
||||
return
|
||||
}
|
||||
}
|
||||
if tc.Function.Arguments != "" {
|
||||
state.arguments += tc.Function.Arguments
|
||||
if !writeEvent("response.function_call_arguments.delta", map[string]interface{}{
|
||||
"item_id": state.itemID,
|
||||
"output_index": toolOutputIndex,
|
||||
"delta": tc.Function.Arguments,
|
||||
}) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
state, found := outputs[outputIndex]
|
||||
if !found {
|
||||
itemID = fmt.Sprintf("msg_%s", uuid.New().String()[:8])
|
||||
itemIDs[outputIndex] = itemID
|
||||
state = &outputItemState{
|
||||
itemType: "message",
|
||||
itemID: fmt.Sprintf("msg_%s", uuid.New().String()[:8]),
|
||||
}
|
||||
outputs[outputIndex] = state
|
||||
if !writeEvent("response.output_item.added", map[string]interface{}{
|
||||
"output_index": outputIndex,
|
||||
"item": map[string]interface{}{
|
||||
"type": "message",
|
||||
"id": itemID,
|
||||
"id": state.itemID,
|
||||
"role": "assistant",
|
||||
"status": "in_progress",
|
||||
"content": []interface{}{},
|
||||
@@ -662,7 +762,7 @@ func writeResponsesAPIStreaming(t testing.TB, w http.ResponseWriter, r *http.Req
|
||||
return
|
||||
}
|
||||
if !writeEvent("response.content_part.added", map[string]interface{}{
|
||||
"item_id": itemID,
|
||||
"item_id": state.itemID,
|
||||
"output_index": outputIndex,
|
||||
"content_index": 0,
|
||||
"part": map[string]interface{}{
|
||||
@@ -674,9 +774,9 @@ func writeResponsesAPIStreaming(t testing.TB, w http.ResponseWriter, r *http.Req
|
||||
}
|
||||
}
|
||||
|
||||
itemTexts[outputIndex] += choice.Delta
|
||||
state.text += choice.Delta
|
||||
if !writeEvent("response.output_text.delta", map[string]interface{}{
|
||||
"item_id": itemID,
|
||||
"item_id": state.itemID,
|
||||
"output_index": outputIndex,
|
||||
"content_index": 0,
|
||||
"delta": choice.Delta,
|
||||
|
||||
@@ -237,6 +237,63 @@ func TestOpenAI_ToolCalls(t *testing.T) {
|
||||
require.GreaterOrEqual(t, requestCount.Load(), int32(2), "expected follow-up model call after tool execution")
|
||||
}
|
||||
|
||||
func TestOpenAI_ToolCalls_ResponsesAPI(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
var requestCount atomic.Int32
|
||||
serverURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse {
|
||||
switch requestCount.Add(1) {
|
||||
case 1:
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAIToolCallChunk("get_weather", `{"location":"San Francisco"}`),
|
||||
)
|
||||
default:
|
||||
return chattest.OpenAIStreamingResponse(
|
||||
chattest.OpenAITextChunks("The weather in San Francisco is 72F.")...,
|
||||
)
|
||||
}
|
||||
})
|
||||
|
||||
client, err := fantasyopenai.New(
|
||||
fantasyopenai.WithAPIKey("test-key"),
|
||||
fantasyopenai.WithBaseURL(serverURL),
|
||||
fantasyopenai.WithUseResponsesAPI(),
|
||||
)
|
||||
require.NoError(t, err)
|
||||
|
||||
ctx := context.Background()
|
||||
model, err := client.LanguageModel(ctx, "gpt-4")
|
||||
require.NoError(t, err)
|
||||
|
||||
type weatherInput struct {
|
||||
Location string `json:"location"`
|
||||
}
|
||||
var toolCallCount atomic.Int32
|
||||
weatherTool := fantasy.NewAgentTool(
|
||||
"get_weather",
|
||||
"Get weather for a location.",
|
||||
func(ctx context.Context, input weatherInput, _ fantasy.ToolCall) (fantasy.ToolResponse, error) {
|
||||
toolCallCount.Add(1)
|
||||
require.Equal(t, "San Francisco", input.Location)
|
||||
return fantasy.NewTextResponse("72F"), nil
|
||||
},
|
||||
)
|
||||
|
||||
agent := fantasy.NewAgent(
|
||||
model,
|
||||
fantasy.WithSystemPrompt("You are a helpful assistant."),
|
||||
fantasy.WithTools(weatherTool),
|
||||
)
|
||||
|
||||
result, err := agent.Stream(ctx, fantasy.AgentStreamCall{
|
||||
Prompt: "What's the weather in San Francisco?",
|
||||
})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, result)
|
||||
require.Equal(t, int32(1), toolCallCount.Load(), "expected exactly one tool execution")
|
||||
require.GreaterOrEqual(t, requestCount.Load(), int32(2), "expected follow-up model call after tool execution")
|
||||
}
|
||||
|
||||
func TestOpenAI_NonStreaming_ResponsesAPI(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
@@ -3,6 +3,7 @@ package chatd
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"slices"
|
||||
@@ -14,13 +15,15 @@ import (
|
||||
"tailscale.com/util/singleflight"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
|
||||
const (
|
||||
chatConfigProvidersTTL = 10 * time.Second
|
||||
chatConfigModelConfigTTL = 10 * time.Second
|
||||
chatConfigUserPromptTTL = 5 * time.Second
|
||||
chatConfigProvidersTTL = 10 * time.Second
|
||||
chatConfigModelConfigTTL = 10 * time.Second
|
||||
chatConfigUserPromptTTL = 5 * time.Second
|
||||
chatConfigAdvisorConfigTTL = 10 * time.Second
|
||||
// Bound user-prompt cache cardinality so one-shot users do not
|
||||
// accumulate forever in long-lived chatd processes.
|
||||
chatConfigUserPromptEntryLimit = 64 * 1024
|
||||
@@ -31,6 +34,11 @@ type cachedProviders struct {
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
type cachedAdvisorConfig struct {
|
||||
config codersdk.AdvisorConfig
|
||||
expiresAt time.Time
|
||||
}
|
||||
|
||||
type cachedModelConfig struct {
|
||||
config database.ChatModelConfig
|
||||
expiresAt time.Time
|
||||
@@ -82,6 +90,11 @@ type chatConfigCache struct {
|
||||
userPromptEpoch uint64
|
||||
userPrompts *tlru.Cache[uuid.UUID, string]
|
||||
userPromptFetches singleflight.Group[string, string]
|
||||
|
||||
// Advisor configuration (singleton).
|
||||
advisorConfig *cachedAdvisorConfig
|
||||
advisorConfigGeneration uint64
|
||||
advisorConfigFetches singleflight.Group[string, codersdk.AdvisorConfig]
|
||||
}
|
||||
|
||||
func newChatConfigCache(ctx context.Context, db database.Store, clock quartz.Clock) *chatConfigCache {
|
||||
@@ -410,3 +423,97 @@ func (c *chatConfigCache) InvalidateUserPrompt(userID uuid.UUID) {
|
||||
c.userPromptEpoch++
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
// InvalidateAdvisorConfig drops the cached advisor configuration so the
|
||||
// next AdvisorConfig call re-fetches from the database. Called from the
|
||||
// ChatConfigEvent subscriber after an admin writes
|
||||
// PUT /api/experimental/chats/config/advisor; without this the cache
|
||||
// could serve stale enabled/model/limits for up to
|
||||
// chatConfigAdvisorConfigTTL. Bumping the generation counter also
|
||||
// discards any in-flight fill started before the invalidation, so a
|
||||
// stale DB read cannot re-cache the pre-update value.
|
||||
func (c *chatConfigCache) InvalidateAdvisorConfig() {
|
||||
c.mu.Lock()
|
||||
c.advisorConfig = nil
|
||||
c.advisorConfigGeneration++
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
// AdvisorConfig returns the deployment-wide advisor configuration. The
|
||||
// underlying site-config row changes on the order of hours or days, so
|
||||
// this cache saves a per-turn DB round trip on chats that reference the
|
||||
// advisor. Parse errors and lookup errors are surfaced to the caller;
|
||||
// callers that prefer silent fallback handle that at the call site.
|
||||
func (c *chatConfigCache) AdvisorConfig(ctx context.Context) (codersdk.AdvisorConfig, error) {
|
||||
if config, ok := c.cachedAdvisorConfig(); ok {
|
||||
return config, nil
|
||||
}
|
||||
|
||||
generation := c.advisorConfigGenerationSnapshot()
|
||||
config, err := singleflightDoChan(
|
||||
ctx,
|
||||
&c.advisorConfigFetches,
|
||||
fmt.Sprintf("%d:advisor", generation),
|
||||
func() (codersdk.AdvisorConfig, error) {
|
||||
if cached, ok := c.cachedAdvisorConfig(); ok {
|
||||
return cached, nil
|
||||
}
|
||||
|
||||
raw, err := c.db.GetChatAdvisorConfig(c.ctx)
|
||||
if err != nil {
|
||||
return codersdk.AdvisorConfig{}, err
|
||||
}
|
||||
var cfg codersdk.AdvisorConfig
|
||||
if err := json.Unmarshal([]byte(raw), &cfg); err != nil {
|
||||
return codersdk.AdvisorConfig{}, err
|
||||
}
|
||||
c.storeAdvisorConfig(generation, cfg)
|
||||
return cfg, nil
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
return codersdk.AdvisorConfig{}, err
|
||||
}
|
||||
return config, nil
|
||||
}
|
||||
|
||||
func (c *chatConfigCache) cachedAdvisorConfig() (codersdk.AdvisorConfig, bool) {
|
||||
c.mu.RLock()
|
||||
entry := c.advisorConfig
|
||||
c.mu.RUnlock()
|
||||
if entry == nil {
|
||||
return codersdk.AdvisorConfig{}, false
|
||||
}
|
||||
if c.clock.Now().Before(entry.expiresAt) {
|
||||
return entry.config, true
|
||||
}
|
||||
|
||||
c.mu.Lock()
|
||||
if current := c.advisorConfig; current != nil && !c.clock.Now().Before(current.expiresAt) {
|
||||
c.advisorConfig = nil
|
||||
}
|
||||
c.mu.Unlock()
|
||||
|
||||
return codersdk.AdvisorConfig{}, false
|
||||
}
|
||||
|
||||
func (c *chatConfigCache) advisorConfigGenerationSnapshot() uint64 {
|
||||
c.mu.RLock()
|
||||
generation := c.advisorConfigGeneration
|
||||
c.mu.RUnlock()
|
||||
return generation
|
||||
}
|
||||
|
||||
func (c *chatConfigCache) storeAdvisorConfig(generation uint64, config codersdk.AdvisorConfig) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
|
||||
if c.advisorConfigGeneration != generation {
|
||||
return
|
||||
}
|
||||
|
||||
c.advisorConfig = &cachedAdvisorConfig{
|
||||
config: config,
|
||||
expiresAt: c.clock.Now().Add(chatConfigAdvisorConfigTTL),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"golang.org/x/xerrors"
|
||||
|
||||
"github.com/coder/coder/v2/coderd/database"
|
||||
"github.com/coder/coder/v2/codersdk"
|
||||
"github.com/coder/coder/v2/testutil"
|
||||
"github.com/coder/quartz"
|
||||
)
|
||||
@@ -25,11 +26,13 @@ type stubChatConfigStore struct {
|
||||
getChatModelConfigByID func(context.Context, uuid.UUID) (database.ChatModelConfig, error)
|
||||
getDefaultChatModelConfig func(context.Context) (database.ChatModelConfig, error)
|
||||
getUserChatCustomPrompt func(context.Context, uuid.UUID) (string, error)
|
||||
getChatAdvisorConfig func(context.Context) (string, error)
|
||||
|
||||
enabledProvidersCalls atomic.Int32
|
||||
modelConfigByIDCalls atomic.Int32
|
||||
defaultModelConfigCall atomic.Int32
|
||||
userPromptCalls atomic.Int32
|
||||
advisorConfigCalls atomic.Int32
|
||||
}
|
||||
|
||||
func (s *stubChatConfigStore) GetEnabledChatProviders(ctx context.Context) ([]database.ChatProvider, error) {
|
||||
@@ -64,6 +67,14 @@ func (s *stubChatConfigStore) GetUserChatCustomPrompt(ctx context.Context, userI
|
||||
return s.getUserChatCustomPrompt(ctx, userID)
|
||||
}
|
||||
|
||||
func (s *stubChatConfigStore) GetChatAdvisorConfig(ctx context.Context) (string, error) {
|
||||
s.advisorConfigCalls.Add(1)
|
||||
if s.getChatAdvisorConfig == nil {
|
||||
panic("unexpected GetChatAdvisorConfig call")
|
||||
}
|
||||
return s.getChatAdvisorConfig(ctx)
|
||||
}
|
||||
|
||||
func TestConfigCache_EnabledProviders_CacheHit(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
@@ -976,3 +987,216 @@ func TestConfigCache_CallerCancellation(t *testing.T) {
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestConfigCache_AdvisorConfig_CacheHit(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
clock := quartz.NewMock(t)
|
||||
const raw = `{"enabled":true,"max_uses_per_run":3,"max_output_tokens":16384}`
|
||||
store := &stubChatConfigStore{
|
||||
getChatAdvisorConfig: func(context.Context) (string, error) {
|
||||
return raw, nil
|
||||
},
|
||||
}
|
||||
cache := newChatConfigCache(ctx, store, clock)
|
||||
|
||||
first, err := cache.AdvisorConfig(ctx)
|
||||
require.NoError(t, err)
|
||||
second, err := cache.AdvisorConfig(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.True(t, first.Enabled)
|
||||
require.Equal(t, 3, first.MaxUsesPerRun)
|
||||
require.Equal(t, int64(16384), first.MaxOutputTokens)
|
||||
require.Equal(t, first, second)
|
||||
require.Equal(t, int32(1), store.advisorConfigCalls.Load(),
|
||||
"second lookup must be served from cache")
|
||||
}
|
||||
|
||||
func TestConfigCache_AdvisorConfig_TTLExpiry(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
clock := quartz.NewMock(t)
|
||||
store := &stubChatConfigStore{}
|
||||
store.getChatAdvisorConfig = func(context.Context) (string, error) {
|
||||
call := store.advisorConfigCalls.Load()
|
||||
return fmt.Sprintf(`{"max_uses_per_run":%d}`, call), nil
|
||||
}
|
||||
cache := newChatConfigCache(ctx, store, clock)
|
||||
|
||||
first, err := cache.AdvisorConfig(ctx)
|
||||
require.NoError(t, err)
|
||||
clock.Advance(chatConfigAdvisorConfigTTL).MustWait(ctx)
|
||||
second, err := cache.AdvisorConfig(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NotEqual(t, first.MaxUsesPerRun, second.MaxUsesPerRun,
|
||||
"TTL expiry must trigger a refetch")
|
||||
require.Equal(t, int32(2), store.advisorConfigCalls.Load())
|
||||
}
|
||||
|
||||
func TestConfigCache_AdvisorConfig_DBErrorNotCached(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
clock := quartz.NewMock(t)
|
||||
expected := xerrors.New("boom")
|
||||
store := &stubChatConfigStore{
|
||||
getChatAdvisorConfig: func(context.Context) (string, error) {
|
||||
return "", expected
|
||||
},
|
||||
}
|
||||
cache := newChatConfigCache(ctx, store, clock)
|
||||
|
||||
_, err := cache.AdvisorConfig(ctx)
|
||||
require.ErrorIs(t, err, expected)
|
||||
_, err = cache.AdvisorConfig(ctx)
|
||||
require.ErrorIs(t, err, expected)
|
||||
|
||||
require.Equal(t, int32(2), store.advisorConfigCalls.Load(),
|
||||
"errors must not populate the cache; every call retries")
|
||||
}
|
||||
|
||||
func TestConfigCache_AdvisorConfig_InvalidJSONNotCached(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
clock := quartz.NewMock(t)
|
||||
store := &stubChatConfigStore{
|
||||
getChatAdvisorConfig: func(context.Context) (string, error) {
|
||||
return "not valid json", nil
|
||||
},
|
||||
}
|
||||
cache := newChatConfigCache(ctx, store, clock)
|
||||
|
||||
_, err := cache.AdvisorConfig(ctx)
|
||||
require.Error(t, err, "malformed JSON must surface as an error")
|
||||
_, err = cache.AdvisorConfig(ctx)
|
||||
require.Error(t, err)
|
||||
|
||||
require.Equal(t, int32(2), store.advisorConfigCalls.Load(),
|
||||
"parse errors must not populate the cache; every call retries")
|
||||
}
|
||||
|
||||
func TestConfigCache_AdvisorConfig_EmptyJSONYieldsZeroValue(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
clock := quartz.NewMock(t)
|
||||
// GetChatAdvisorConfig returns "{}" when the site-config row is
|
||||
// absent. That must unmarshal to a zero-value AdvisorConfig rather
|
||||
// than a parse error.
|
||||
store := &stubChatConfigStore{
|
||||
getChatAdvisorConfig: func(context.Context) (string, error) {
|
||||
return "{}", nil
|
||||
},
|
||||
}
|
||||
cache := newChatConfigCache(ctx, store, clock)
|
||||
|
||||
cfg, err := cache.AdvisorConfig(ctx)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, codersdk.AdvisorConfig{}, cfg)
|
||||
}
|
||||
|
||||
// Guards the pubsub-driven invalidation path. Without this, an admin
|
||||
// writing PUT /api/experimental/chats/config/advisor could keep every
|
||||
// replica serving stale enabled/model/limits for up to
|
||||
// chatConfigAdvisorConfigTTL, which defeats the subscriber in chatd.go.
|
||||
func TestConfigCache_InvalidateAdvisorConfig(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
clock := quartz.NewMock(t)
|
||||
store := &stubChatConfigStore{}
|
||||
store.getChatAdvisorConfig = func(context.Context) (string, error) {
|
||||
call := store.advisorConfigCalls.Load()
|
||||
return fmt.Sprintf(`{"max_uses_per_run":%d}`, call), nil
|
||||
}
|
||||
cache := newChatConfigCache(ctx, store, clock)
|
||||
|
||||
first, err := cache.AdvisorConfig(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
cache.InvalidateAdvisorConfig()
|
||||
|
||||
second, err := cache.AdvisorConfig(ctx)
|
||||
require.NoError(t, err)
|
||||
|
||||
require.NotEqual(t, first.MaxUsesPerRun, second.MaxUsesPerRun,
|
||||
"invalidation must force a refetch without waiting for TTL expiry")
|
||||
require.Equal(t, int32(2), store.advisorConfigCalls.Load())
|
||||
}
|
||||
|
||||
// Guards against the invalidation-during-singleflight race. A stale
|
||||
// in-flight fill started before InvalidateAdvisorConfig must not
|
||||
// re-cache its pre-update value, which would defeat the pubsub
|
||||
// invalidation path for up to chatConfigAdvisorConfigTTL.
|
||||
func TestConfigCache_InvalidateAdvisorConfig_BlocksStaleInFlight(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ctx := testutil.Context(t, testutil.WaitMedium)
|
||||
clock := quartz.NewMock(t)
|
||||
staleConfig := `{"max_uses_per_run":1}`
|
||||
freshConfig := `{"max_uses_per_run":2}`
|
||||
firstStarted := make(chan struct{})
|
||||
secondStarted := make(chan struct{})
|
||||
releaseFirst := make(chan struct{})
|
||||
releaseSecond := make(chan struct{})
|
||||
store := &stubChatConfigStore{}
|
||||
store.getChatAdvisorConfig = func(context.Context) (string, error) {
|
||||
switch call := store.advisorConfigCalls.Load(); call {
|
||||
case 1:
|
||||
close(firstStarted)
|
||||
<-releaseFirst
|
||||
return staleConfig, nil
|
||||
case 2:
|
||||
close(secondStarted)
|
||||
<-releaseSecond
|
||||
return freshConfig, nil
|
||||
default:
|
||||
return "", xerrors.Errorf("unexpected advisor config call %d", call)
|
||||
}
|
||||
}
|
||||
cache := newChatConfigCache(ctx, store, clock)
|
||||
|
||||
type result struct {
|
||||
config codersdk.AdvisorConfig
|
||||
err error
|
||||
}
|
||||
|
||||
firstResult := make(chan result, 1)
|
||||
go func() {
|
||||
config, err := cache.AdvisorConfig(ctx)
|
||||
firstResult <- result{config: config, err: err}
|
||||
}()
|
||||
|
||||
waitForSignal(t, firstStarted)
|
||||
cache.InvalidateAdvisorConfig()
|
||||
|
||||
secondResult := make(chan result, 1)
|
||||
go func() {
|
||||
config, err := cache.AdvisorConfig(ctx)
|
||||
secondResult <- result{config: config, err: err}
|
||||
}()
|
||||
|
||||
waitForSignal(t, secondStarted)
|
||||
close(releaseFirst)
|
||||
first := <-firstResult
|
||||
require.NoError(t, first.err)
|
||||
require.EqualValues(t, 1, first.config.MaxUsesPerRun)
|
||||
require.Nil(t, cache.advisorConfig,
|
||||
"stale fill must not re-cache after invalidation")
|
||||
|
||||
close(releaseSecond)
|
||||
second := <-secondResult
|
||||
require.NoError(t, second.err)
|
||||
require.EqualValues(t, 2, second.config.MaxUsesPerRun)
|
||||
require.Equal(t, int32(2), store.advisorConfigCalls.Load())
|
||||
|
||||
third, err := cache.AdvisorConfig(ctx)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 2, third.MaxUsesPerRun)
|
||||
require.Equal(t, int32(2), store.advisorConfigCalls.Load())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user