From b4809a6b768f19eb60c240bd5ba0df63968e86a3 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Sun, 12 Jul 2026 21:35:51 +0000 Subject: [PATCH 1/2] feat: gate chat turns on reported workspace context A workspace-bound chat turn no longer runs before the workspace agent has reported its context snapshot, and the first snapshot now already contains the MCP catalog: - agent: connect workspace MCP servers before releasing the context gate, bounded by a 30s first-sync wait, so the initial push carries MCP servers and tools. Late-settling servers still re-push via the catalog-change callback. Lifecycle Ready is unaffected. - agentcontext: include MCP config and live MCP server state in the aggregate hash, so MCP changes surface as chat-context drift instead of being silently missed by pinned chats. - chatd: treat the context report as a hard precondition for workspace turns. Turn prep rebinds stale agent bindings to the latest start build (re-pinning context), waits on a 1s poll up to a 15m ceiling for the agent's first report, and fails fast with actionable errors when the workspace is not started or the agent predates Agent API 2.10. Gate failures surface as the chat's visible error state and interrupts keep working during the wait. - surface: ChatContext gains a waiting/ready state, first hydration and gate-exit pinning publish a context_ready watch event, and the agents UI renders the waiting state. --- agent/agent.go | 33 +- agent/agent_context_test.go | 147 ++++ agent/agentcontext/mcp_internal_test.go | 54 ++ agent/agentcontext/resolve.go | 42 +- agent/agentcontext/resolve_test.go | 77 ++- agent/x/agentmcp/manager.go | 18 + agent/x/agentmcp/manager_internal_test.go | 92 +++ coderd/apidoc/docs.go | 23 +- coderd/apidoc/swagger.json | 17 +- coderd/database/db2sdk/db2sdk.go | 15 +- coderd/database/db2sdk/db2sdk_test.go | 87 +++ coderd/database/dbauthz/dbauthz.go | 4 +- coderd/database/dbmetrics/querymetrics.go | 6 +- coderd/database/dbmock/dbmock.go | 7 +- coderd/database/querier.go | 8 +- coderd/database/querier_test.go | 6 +- coderd/database/queries.sql.go | 79 ++- coderd/database/queries/chats.sql | 47 +- coderd/x/chatd/ARCHITECTURE.md | 9 + coderd/x/chatd/chatd.go | 67 +- coderd/x/chatd/chatd_internal_test.go | 32 +- coderd/x/chatd/chatd_test.go | 126 ++-- coderd/x/chatd/context_gate.go | 296 ++++++++ coderd/x/chatd/context_gate_internal_test.go | 640 ++++++++++++++++++ coderd/x/chatd/context_hydration.go | 45 +- .../chatd/context_hydration_internal_test.go | 58 +- coderd/x/chatd/context_integration_test.go | 10 +- .../x/chatd/context_prompt_internal_test.go | 5 +- .../x/chatd/context_rebind_internal_test.go | 15 +- coderd/x/chatd/generation_preparer.go | 30 +- coderd/x/chatd/subagent_internal_test.go | 12 + codersdk/chats.go | 23 +- docs/reference/api/chats.md | 46 +- docs/reference/api/schemas.md | 35 +- site/src/api/queries/chats.test.ts | 43 ++ site/src/api/queries/chats.ts | 14 +- site/src/api/typesGenerated.ts | 15 +- site/src/pages/AgentsPage/AgentsPage.tsx | 5 +- .../ContextUsageIndicator.stories.tsx | 26 + .../components/ContextUsageIndicator.tsx | 28 +- site/src/testHelpers/chatEntities.ts | 9 + 41 files changed, 2044 insertions(+), 307 deletions(-) create mode 100644 coderd/x/chatd/context_gate.go create mode 100644 coderd/x/chatd/context_gate_internal_test.go diff --git a/agent/agent.go b/agent/agent.go index 550392802df..db53193a5f6 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -80,6 +80,12 @@ const ( EnvProcOOMScore = "CODER_PROC_OOM_SCORE" ) +// mcpFirstSyncTimeout bounds how long the startup goroutine waits for +// the first MCP catalog sync before releasing the context gate. The +// reload keeps running in the background on timeout, and late-settling +// servers still re-push via the catalog-change callback. +const mcpFirstSyncTimeout = 30 * time.Second + var ErrAgentClosing = xerrors.New("agent is closing") type Options struct { @@ -1552,19 +1558,24 @@ func (a *agent) handleManifest(manifestOK *checkpoint) func(ctx context.Context, a.metrics.startupScriptSeconds.WithLabelValues(label).Set(dur) a.scriptRunner.StartCron() - // Startup finished (success or terminal failure): release - // the context gate. MCP servers connect below and - // re-trigger a push once up, so we don't block readiness - // on them. - a.contextManager.SetReady() - - // Connect to workspace MCP servers after the - // lifecycle transition to avoid delaying Ready. - // This runs inside the tracked goroutine so it - // is properly awaited on shutdown. - if mcpErr := a.mcpManager.Reload(a.gracefulCtx, a.contextConfigAPI.MCPConfigFiles()); mcpErr != nil { + // Connect to workspace MCP servers before releasing the + // context gate so the first context push already contains + // the MCP catalog. The wait is bounded: if the catalog has + // not settled within mcpFirstSyncTimeout, the gate is + // released anyway while the reload keeps running in the + // background. Lifecycle Ready was reported above and is + // not delayed by this wait. This runs inside the tracked + // goroutine so it is properly awaited on shutdown. + if mcpErr := a.mcpManager.ReloadWithTimeout(a.gracefulCtx, a.contextConfigAPI.MCPConfigFiles(), mcpFirstSyncTimeout); mcpErr != nil { a.logger.Warn(ctx, "failed to reload workspace MCP servers", slog.Error(mcpErr)) } + + // Startup finished (success or terminal failure): release + // the context gate so the first context snapshot resolves + // and is pushed. Servers that settle after the bounded + // wait above still re-trigger a push via the MCP manager's + // catalog-change callback. + a.contextManager.SetReady() }) if err != nil { return xerrors.Errorf("track conn goroutine: %w", err) diff --git a/agent/agent_context_test.go b/agent/agent_context_test.go index c154bb00ed8..3ff50bf252e 100644 --- a/agent/agent_context_test.go +++ b/agent/agent_context_test.go @@ -1,6 +1,9 @@ package agent_test import ( + "bufio" + "encoding/json" + "fmt" "os" "path/filepath" "testing" @@ -68,3 +71,147 @@ func TestAgent_ContextStatePushed(t *testing.T) { assert.False(t, p.GetInitial(), "only the first push must be Initial") } } + +// TestAgent_ContextStateFirstPushIncludesMCP verifies the context gate +// waits for the MCP catalog to settle before the first push: the +// initial PushContextState (Initial=true) already contains the +// configured MCP server resource with its tools instead of relying on +// a follow-up push after the servers connect. +func TestAgent_ContextStateFirstPushIncludesMCP(t *testing.T) { + t.Parallel() + + if os.Getenv("TEST_MCP_FAKE_SERVER") == "1" { + // Child process: act as a minimal MCP server over stdio. + runFakeMCPStdioServer() + return + } + + dir := t.TempDir() + + // Point the default .mcp.json at a fake MCP server using the + // test binary re-exec pattern: the agent spawns this test + // binary, which detects TEST_MCP_FAKE_SERVER and serves MCP + // over stdio instead of running tests. + testBin, err := os.Executable() + require.NoError(t, err) + mcpConfig := map[string]any{ + "mcpServers": map[string]any{ + "srv": map[string]any{ + "command": testBin, + "args": []string{"-test.run=^TestAgent_ContextStateFirstPushIncludesMCP$"}, + "env": map[string]string{"TEST_MCP_FAKE_SERVER": "1"}, + }, + }, + } + data, err := json.Marshal(mcpConfig) + require.NoError(t, err) + require.NoError(t, + os.WriteFile(filepath.Join(dir, ".mcp.json"), data, 0o600)) + + //nolint:dogsled // setupAgent returns a wide tuple; we only care about the client. + _, client, _, _, _ := setupAgent(t, + agentsdk.Manifest{Directory: dir}, + 0, + func(_ *agenttest.Client, opts *agent.Options) { + opts.ContextConfig = agentcontextconfig.Config{} + }, + ) + + var pushes []*agentproto.PushContextStateRequest + require.Eventually(t, func() bool { + pushes = client.ContextStatePushes() + return len(pushes) > 0 + }, testutil.WaitLong, testutil.IntervalFast, + "expected a context snapshot push after startup; got %d pushes", len(pushes)) + + first := pushes[0] + require.True(t, first.GetInitial(), "first push must carry Initial=true") + + // The gate held until the MCP catalog settled, so the first + // push must already contain the server and its tool list. + var srv *agentproto.MCPServerBody + for _, r := range first.GetResources() { + if body := r.GetMcpServer(); body != nil && body.GetServerName() == "srv" { + srv = body + } + } + require.NotNil(t, srv, + "first push must already contain the MCP server resource") + require.Len(t, srv.GetTools(), 1, + "first push must already contain the server's tools") + assert.Equal(t, "echo", srv.GetTools()[0].GetName()) +} + +// runFakeMCPStdioServer implements a minimal JSON-RPC / MCP server +// over stdin/stdout, just enough for initialize + tools/list. It +// runs in a re-exec'd copy of the test binary spawned by the agent's +// MCP manager. +func runFakeMCPStdioServer() { + scanner := bufio.NewScanner(os.Stdin) + for scanner.Scan() { + line := scanner.Bytes() + + var req struct { + JSONRPC string `json:"jsonrpc"` + ID json.RawMessage `json:"id"` + Method string `json:"method"` + } + if err := json.Unmarshal(line, &req); err != nil { + continue + } + + var resp any + switch req.Method { + case "initialize": + resp = map[string]any{ + "jsonrpc": "2.0", + "id": req.ID, + "result": map[string]any{ + "protocolVersion": "2025-03-26", + "capabilities": map[string]any{ + "tools": map[string]any{}, + }, + "serverInfo": map[string]any{ + "name": "fake-server", + "version": "0.0.1", + }, + }, + } + case "notifications/initialized": + // No response needed for notifications. + continue + case "tools/list": + resp = map[string]any{ + "jsonrpc": "2.0", + "id": req.ID, + "result": map[string]any{ + "tools": []map[string]any{ + { + "name": "echo", + "description": "echoes input", + "inputSchema": map[string]any{ + "type": "object", + "properties": map[string]any{}, + }, + }, + }, + }, + } + default: + resp = map[string]any{ + "jsonrpc": "2.0", + "id": req.ID, + "error": map[string]any{ + "code": -32601, + "message": "method not found", + }, + } + } + + out, err := json.Marshal(resp) + if err != nil { + continue + } + _, _ = fmt.Fprintf(os.Stdout, "%s\n", out) + } +} diff --git a/agent/agentcontext/mcp_internal_test.go b/agent/agentcontext/mcp_internal_test.go index 9f101d0b3e7..b1ae584f1ab 100644 --- a/agent/agentcontext/mcp_internal_test.go +++ b/agent/agentcontext/mcp_internal_test.go @@ -152,3 +152,57 @@ func TestBuildMCPServerResources(t *testing.T) { require.Equal(t, StatusOK, got[1].Status) }) } + +// TestAggregateHash_MCPServerState verifies that live MCP server state +// participates in the aggregate hash: tool changes and connect-state +// flips change the hash, while identical inputs and tool order +// permutations do not (buildMCPServerResources sorts tools by name +// before hashing). +func TestAggregateHash_MCPServerState(t *testing.T) { + t.Parallel() + + hashOf := func(servers []MCPServerStatus) [32]byte { + return ComputeAggregateHash(buildMCPServerResources(servers)) + } + + connected := []MCPServerStatus{ + {Name: "fs", Connected: true, Tools: []MCPTool{ + {Name: "read", Description: "Read"}, + {Name: "write", Description: "Write"}, + }}, + } + base := hashOf(connected) + + // Identical inputs hash identically. + require.Equal(t, base, hashOf(connected)) + + // Tool order permutations hash identically because tools are + // sorted by name before hashing. + permuted := []MCPServerStatus{ + {Name: "fs", Connected: true, Tools: []MCPTool{ + {Name: "write", Description: "Write"}, + {Name: "read", Description: "Read"}, + }}, + } + require.Equal(t, base, hashOf(permuted)) + + // A tool-list change flips the aggregate hash. + toolAdded := []MCPServerStatus{ + {Name: "fs", Connected: true, Tools: []MCPTool{ + {Name: "read", Description: "Read"}, + {Name: "write", Description: "Write"}, + {Name: "stat", Description: "Stat"}, + }}, + } + require.NotEqual(t, base, hashOf(toolAdded)) + + // A connected-to-failed flip changes both the resource status and + // its content hash, so the aggregate hash flips too. + failed := []MCPServerStatus{ + {Name: "fs", Connected: false, Err: "connection refused"}, + } + require.NotEqual(t, base, hashOf(failed)) + + // A server disappearing entirely flips the aggregate hash. + require.NotEqual(t, base, hashOf(nil)) +} diff --git a/agent/agentcontext/resolve.go b/agent/agentcontext/resolve.go index 92e713783b9..064fa555138 100644 --- a/agent/agentcontext/resolve.go +++ b/agent/agentcontext/resolve.go @@ -184,10 +184,10 @@ func (r *Resolver) ResolveContext(ctx context.Context, roots []ScanRoot) Snapsho payloadBytes += uint64(len(r.Payload)) } - // The drift hash covers only pinned prompt content; MCP resources are - // excluded (see driftResources). Snapshot.Resources still carries the - // full set so MCP servers stay visible in the chat-context snapshot. - hash := ComputeAggregateHash(driftResources(resources)) + // The aggregate hash covers every resource, including MCP config and + // live MCP server state, so MCP changes (server added or removed, tool + // list changed, connect state changed) surface as chat-context drift. + hash := ComputeAggregateHash(resources) snap := Snapshot{ Resources: resources, @@ -959,12 +959,12 @@ type Snapshot struct { // loop withholds. Version uint64 // AggregateHash is sha256 over a canonical encoding of - // (ID, Kind, Source, ContentHash, Status) for every - // drift-relevant resource. MCP resources (KindMCPConfig and - // KindMCPServer) are excluded because they describe live, - // agent-global runtime capabilities discovered at turn time, - // not pinned prompt content; see driftResources. Identical - // inputs always produce identical hashes; see + // (ID, Kind, Source, ContentHash, Status) for every resource in + // the snapshot, including MCP config and live MCP server state. + // coderd pins MCP servers and their tools into chat context, so + // an MCP change (server added or removed, tool list changed, + // connect state changed) surfaces as chat-context drift. + // Identical inputs always produce identical hashes; see // ComputeAggregateHash. AggregateHash [32]byte // Resources is sorted by ID for deterministic encoding. @@ -978,28 +978,6 @@ type Snapshot struct { SnapshotError string } -// driftResources returns the subset of resources that participate in -// chat-context drift detection. MCP resources (the .mcp.json config and -// connected MCP servers) are deliberately excluded: an agent connects to -// its MCP servers asynchronously after startup, and the chat model -// discovers their tools live at turn time, not from pinned prompt -// content. Hashing them would dirty an already-hydrated chat the moment -// a server finished connecting, even though nothing the user pinned -// changed. Instruction files and skills, whose content is pinned into -// the chat, stay drift-relevant. -func driftResources(resources []Resource) []Resource { - out := make([]Resource, 0, len(resources)) - for _, r := range resources { - switch r.Kind { - case KindMCPConfig, KindMCPServer: - continue - default: - out = append(out, r) - } - } - return out -} - // ComputeAggregateHash produces the deterministic snapshot // aggregate hash for the supplied resources. The caller does // not need to pre-sort; the function sorts a copy of the slice diff --git a/agent/agentcontext/resolve_test.go b/agent/agentcontext/resolve_test.go index a8ea63e951f..6e8d0e0331f 100644 --- a/agent/agentcontext/resolve_test.go +++ b/agent/agentcontext/resolve_test.go @@ -498,37 +498,76 @@ func TestResolver_MCPResourcesRespectAggregateByteCap(t *testing.T) { require.NotEmpty(t, snap.SnapshotError, "snapshot must surface the cap breach") } -// TestResolver_MCPExcludedFromAggregateHash verifies that MCP resources -// (config and live servers) are carried in the snapshot but excluded -// from the drift/aggregate hash, so an MCP server connecting (or its -// tools changing) does not flip already-hydrated chats to dirty. -func TestResolver_MCPExcludedFromAggregateHash(t *testing.T) { +// TestResolver_MCPIncludedInAggregateHash verifies that MCP resources +// (config and live servers) participate in the aggregate hash. coderd +// pins MCP servers and their tools into chat context, so a server +// appearing, its tools changing, or the .mcp.json config changing must +// all surface as chat-context drift. +func TestResolver_MCPIncludedInAggregateHash(t *testing.T) { t.Parallel() dir := t.TempDir() - // An instruction file provides drift-relevant pinned content. mustWriteFile(t, filepath.Join(dir, "AGENTS.md"), "workspace rules") base := (&agentcontext.Resolver{}).Resolve([]agentcontext.ScanRoot{{Path: dir}}) - mcpRes := agentcontext.Resource{ - ID: "mcp_server:github", - Kind: agentcontext.KindMCPServer, - Source: "github", - Name: "github", - Status: agentcontext.StatusOK, - ContentHash: sha256.Sum256([]byte("tool-list")), - Tools: []agentcontext.MCPTool{{Name: "search"}}, + mcpServer := func(contentHash [32]byte) func() []agentcontext.Resource { + return func() []agentcontext.Resource { + return []agentcontext.Resource{{ + ID: "mcp_server:github", + Kind: agentcontext.KindMCPServer, + Source: "github", + Name: "github", + Status: agentcontext.StatusOK, + ContentHash: contentHash, + Tools: []agentcontext.MCPTool{{Name: "search"}}, + }} + } } + toolsV1 := sha256.Sum256([]byte("tool-list-v1")) + toolsV2 := sha256.Sum256([]byte("tool-list-v2")) + withMCP := (&agentcontext.Resolver{ - MCPResources: func() []agentcontext.Resource { return []agentcontext.Resource{mcpRes} }, + MCPResources: mcpServer(toolsV1), }).Resolve([]agentcontext.ScanRoot{{Path: dir}}) - // The MCP server resource is present in the snapshot... + // The MCP server resource is present in the snapshot and changes + // the aggregate hash relative to a snapshot without it. got := findResource(t, withMCP.Resources, agentcontext.KindMCPServer, "github") require.Len(t, got.Tools, 1) - // ...but does not change the drift/aggregate hash. - require.Equal(t, base.AggregateHash, withMCP.AggregateHash, - "MCP resources must not participate in the drift hash") + require.NotEqual(t, base.AggregateHash, withMCP.AggregateHash, + "an MCP server appearing must change the aggregate hash") + + // Identical inputs produce an identical hash. + sameMCP := (&agentcontext.Resolver{ + MCPResources: mcpServer(toolsV1), + }).Resolve([]agentcontext.ScanRoot{{Path: dir}}) + require.Equal(t, withMCP.AggregateHash, sameMCP.AggregateHash, + "identical inputs must hash identically") + + // A tool-list change (a new content hash) changes the aggregate hash. + changedTools := (&agentcontext.Resolver{ + MCPResources: mcpServer(toolsV2), + }).Resolve([]agentcontext.ScanRoot{{Path: dir}}) + require.NotEqual(t, withMCP.AggregateHash, changedTools.AggregateHash, + "an MCP server's tools changing must change the aggregate hash") +} + +// TestResolver_MCPConfigChangesAggregateHash verifies that an .mcp.json +// config resource participates in the aggregate hash, so editing the +// config surfaces as chat-context drift. +func TestResolver_MCPConfigChangesAggregateHash(t *testing.T) { + t.Parallel() + dir := t.TempDir() + path := filepath.Join(dir, ".mcp.json") + + mustWriteFile(t, path, `{"mcpServers": {}}`) + base := (&agentcontext.Resolver{}).Resolve([]agentcontext.ScanRoot{{Path: dir}}) + findResource(t, base.Resources, agentcontext.KindMCPConfig, path) + + mustWriteFile(t, path, `{"mcpServers": {"github": {"command": "gh-mcp"}}}`) + changed := (&agentcontext.Resolver{}).Resolve([]agentcontext.ScanRoot{{Path: dir}}) + require.NotEqual(t, base.AggregateHash, changed.AggregateHash, + "an MCP config change must change the aggregate hash") } // TestResolver_UnreadableInstructionFile verifies the diff --git a/agent/x/agentmcp/manager.go b/agent/x/agentmcp/manager.go index 363f61a3dfb..22f76da7378 100644 --- a/agent/x/agentmcp/manager.go +++ b/agent/x/agentmcp/manager.go @@ -177,6 +177,24 @@ func (m *Manager) Reload(ctx context.Context, paths []string) error { return m.waitReload(ctx, ch, 0) } +// ReloadWithTimeout behaves like Reload but bounds the caller's wait: +// it waits up to timeout for the reload to settle and returns a +// context.DeadlineExceeded-wrapping error if it does not. The reload +// itself keeps running in the background under the manager context, +// so a timed-out caller does not cancel it and a later Reload or +// catalog callback observes the eventual result. A timeout of zero +// or less waits indefinitely, matching Reload. +func (m *Manager) ReloadWithTimeout(ctx context.Context, paths []string, timeout time.Duration) error { + ch, started, err := m.startReloadIfNeeded(paths) + if err != nil { + return err + } + if !started { + return nil + } + return m.waitReload(ctx, ch, timeout) +} + // SetOnReload registers a callback fired (outside the cache lock) after // a reload changes the per-server catalog. The agent wires this to the // agentcontext manager's Trigger so discovery re-resolves and re-pushes diff --git a/agent/x/agentmcp/manager_internal_test.go b/agent/x/agentmcp/manager_internal_test.go index 8ec8bc77e83..8ec0f68f5b0 100644 --- a/agent/x/agentmcp/manager_internal_test.go +++ b/agent/x/agentmcp/manager_internal_test.go @@ -6,6 +6,7 @@ import ( "encoding/json" "fmt" "os" + "sync" "testing" "time" @@ -277,6 +278,97 @@ func TestManager_WaitReloadTimeout(t *testing.T) { assert.Contains(t, err.Error(), "tools reload timed out after 1m0s") } +// TestReloadWithTimeout_Settles verifies that ReloadWithTimeout +// returns nil once the reload settles within the timeout and the +// catalog reflects the connected server. +func TestReloadWithTimeout_Settles(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug) + dir := t.TempDir() + + _, entry := fakeMCPServerConfig(t, "srv") + configPath := writeMCPConfig(t, dir, map[string]mcpServerEntry{"srv": entry}) + + m := NewManager(ctx, logger, agentexec.DefaultExecer, nil) + t.Cleanup(func() { _ = m.Close() }) + + err := m.ReloadWithTimeout(ctx, []string{configPath}, time.Minute) + require.NoError(t, err) + + tools := m.connectedTools() + require.Len(t, tools, 1) + assert.Equal(t, "echo", tools[0].tool) + + // A fresh snapshot short-circuits the next bounded reload. + err = m.ReloadWithTimeout(ctx, []string{configPath}, time.Minute) + require.NoError(t, err) +} + +// TestReloadWithTimeout_TimesOut verifies that ReloadWithTimeout +// returns a deadline error when the reload has not settled by the +// time the timer fires, and that the reload keeps running in the +// background so the catalog still settles afterwards. +func TestReloadWithTimeout_TimesOut(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + logger := slogtest.Make(t, nil).Leveled(slog.LevelDebug) + clock := quartz.NewMock(t) + timerTrap := clock.Trap().NewTimer("agentmcp", "tools_reload") + defer timerTrap.Close() + dir := t.TempDir() + + _, entry := fakeMCPServerConfig(t, "srv") + configPath := writeMCPConfig(t, dir, map[string]mcpServerEntry{"srv": entry}) + + m := NewManager(ctx, logger, agentexec.DefaultExecer, nil) + m.clock = clock + t.Cleanup(func() { _ = m.Close() }) + + // Block the reload body in connectAll so it cannot settle + // before the timer fires. Release runs on cleanup too, before + // Close, so a failed assertion cannot leave the hook blocked. + reached := make(chan struct{}) + release := make(chan struct{}) + var reachedOnce, releaseOnce sync.Once + releaseHook := func() { releaseOnce.Do(func() { close(release) }) } + t.Cleanup(releaseHook) + m.connectStartedHook = func() { + reachedOnce.Do(func() { close(reached) }) + select { + case <-release: + case <-ctx.Done(): + } + } + + done := make(chan error, 1) + go func() { + done <- m.ReloadWithTimeout(ctx, []string{configPath}, time.Minute) + }() + + // The reload body is blocked inside the hook while the caller + // waits on the bounded timer. + testutil.TryReceive(ctx, t, reached) + call := timerTrap.MustWait(ctx) + require.Equal(t, time.Minute, call.Duration) + call.MustRelease(ctx) + + clock.Advance(time.Minute).MustWait(ctx) + err := testutil.RequireReceive(ctx, t, done) + require.Error(t, err) + assert.ErrorIs(t, err, context.DeadlineExceeded) + + // The reload was not canceled by the timed-out caller: once + // unblocked it settles and the catalog contains the server. + releaseHook() + require.Eventually(t, func() bool { + return len(m.connectedTools()) == 1 + }, testutil.WaitLong, testutil.IntervalFast, + "catalog should settle after the timed-out reload completes") +} + // runFakeMCPServer implements a minimal JSON-RPC / MCP server over // stdin/stdout, just enough for initialize + tools/list. func runFakeMCPServer() { diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 349aea763e9..a5afa4d085d 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -16709,7 +16709,7 @@ const docTemplate = `{ "$ref": "#/definitions/codersdk.ChatClientType" }, "context": { - "description": "Context reports the chat's pinned workspace-context state and\nwhether it has drifted from the agent's latest pushed snapshot.\nNil when the chat has no pinned context yet.", + "description": "Context reports the chat's pinned workspace-context state and\nwhether it has drifted from the agent's latest pushed snapshot.\nPresent for every workspace-bound chat; nil for chats with no\nworkspace and no context remnants.", "allOf": [ { "$ref": "#/definitions/codersdk.ChatContext" @@ -16889,6 +16889,14 @@ const docTemplate = `{ "items": { "$ref": "#/definitions/codersdk.ChatContextResource" } + }, + "state": { + "description": "State reports where the chat stands in the context-reporting\nlifecycle: waiting for the agent's first report or ready (pinned).", + "allOf": [ + { + "$ref": "#/definitions/codersdk.ChatContextState" + } + ] } } }, @@ -16966,6 +16974,17 @@ const docTemplate = `{ "ChatContextResourceStatusExcluded" ] }, + "codersdk.ChatContextState": { + "type": "string", + "enum": [ + "waiting", + "ready" + ], + "x-enum-varnames": [ + "ChatContextStateWaiting", + "ChatContextStateReady" + ] + }, "codersdk.ChatContextTool": { "type": "object", "properties": { @@ -17938,6 +17957,7 @@ const docTemplate = `{ "deleted", "diff_status_change", "action_required", + "context_ready", "context_dirty" ], "x-enum-varnames": [ @@ -17948,6 +17968,7 @@ const docTemplate = `{ "ChatWatchEventKindDeleted", "ChatWatchEventKindDiffStatusChange", "ChatWatchEventKindActionRequired", + "ChatWatchEventKindContextReady", "ChatWatchEventKindContextDirty" ] }, diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 2e24bfa51e2..f1455a26b3b 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -15006,7 +15006,7 @@ "$ref": "#/definitions/codersdk.ChatClientType" }, "context": { - "description": "Context reports the chat's pinned workspace-context state and\nwhether it has drifted from the agent's latest pushed snapshot.\nNil when the chat has no pinned context yet.", + "description": "Context reports the chat's pinned workspace-context state and\nwhether it has drifted from the agent's latest pushed snapshot.\nPresent for every workspace-bound chat; nil for chats with no\nworkspace and no context remnants.", "allOf": [ { "$ref": "#/definitions/codersdk.ChatContext" @@ -15174,6 +15174,14 @@ "items": { "$ref": "#/definitions/codersdk.ChatContextResource" } + }, + "state": { + "description": "State reports where the chat stands in the context-reporting\nlifecycle: waiting for the agent's first report or ready (pinned).", + "allOf": [ + { + "$ref": "#/definitions/codersdk.ChatContextState" + } + ] } } }, @@ -15240,6 +15248,11 @@ "ChatContextResourceStatusExcluded" ] }, + "codersdk.ChatContextState": { + "type": "string", + "enum": ["waiting", "ready"], + "x-enum-varnames": ["ChatContextStateWaiting", "ChatContextStateReady"] + }, "codersdk.ChatContextTool": { "type": "object", "properties": { @@ -16182,6 +16195,7 @@ "deleted", "diff_status_change", "action_required", + "context_ready", "context_dirty" ], "x-enum-varnames": [ @@ -16192,6 +16206,7 @@ "ChatWatchEventKindDeleted", "ChatWatchEventKindDiffStatusChange", "ChatWatchEventKindActionRequired", + "ChatWatchEventKindContextReady", "ChatWatchEventKindContextDirty" ] }, diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index 4754bbbe950..86aedbe4ae6 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1766,13 +1766,22 @@ func Chat(c database.Chat, diffStatus *database.ChatDiffStatus, files []database }) } } - // Report pinned-context state when the chat is context-tracked - // (has a pinned hash), dirty, or carries a snapshot error. - if len(c.ContextAggregateHash) > 0 || c.ContextDirtySince.Valid || c.ContextError != "" { + // Report pinned-context state when the chat is workspace-bound or + // context-tracked (has a pinned hash, is dirty, or carries a snapshot + // error). Workspace-bound chats always get a Context so the client can + // distinguish waiting (not yet pinned to a reported snapshot) from + // ready (pinned). + if c.WorkspaceID.Valid || len(c.ContextAggregateHash) > 0 || c.ContextDirtySince.Valid || c.ContextError != "" { chatContext := &codersdk.ChatContext{ Dirty: c.ContextDirtySince.Valid, Error: c.ContextError, } + switch { + case len(c.ContextAggregateHash) > 0: + chatContext.State = codersdk.ChatContextStateReady + case c.WorkspaceID.Valid: + chatContext.State = codersdk.ChatContextStateWaiting + } if c.ContextDirtySince.Valid { dirtySince := c.ContextDirtySince.Time chatContext.DirtySince = &dirtySince diff --git a/coderd/database/db2sdk/db2sdk_test.go b/coderd/database/db2sdk/db2sdk_test.go index 44d43442d20..89c9c42f7dd 100644 --- a/coderd/database/db2sdk/db2sdk_test.go +++ b/coderd/database/db2sdk/db2sdk_test.go @@ -872,6 +872,93 @@ func TestChat_FileMetadataConversion(t *testing.T) { require.NotContains(t, string(data), `"mimetype"`) } +// TestChat_ContextState covers when the converter emits a Context and which +// State it reports: workspace-bound chats always carry one (ready when a +// pinned hash exists, waiting otherwise), non-workspace chats carry one only +// when dirty/error remnants exist (with no state), and plain non-workspace +// chats keep a nil Context. +func TestChat_ContextState(t *testing.T) { + t.Parallel() + + now := dbtime.Now() + workspaceID := uuid.NullUUID{UUID: uuid.New(), Valid: true} + + tests := []struct { + name string + chat database.Chat + wantContext bool + wantState codersdk.ChatContextState + wantDirty bool + }{ + { + name: "WorkspaceBoundPinnedIsReady", + chat: database.Chat{ + WorkspaceID: workspaceID, + ContextAggregateHash: []byte{0x01}, + }, + wantContext: true, + wantState: codersdk.ChatContextStateReady, + }, + { + name: "WorkspaceBoundUnpinnedIsWaiting", + chat: database.Chat{ + WorkspaceID: workspaceID, + }, + wantContext: true, + wantState: codersdk.ChatContextStateWaiting, + }, + { + // A rebind clears the pin to an empty non-NULL hash; the chat + // waits again for the new agent's report. + name: "WorkspaceBoundRebindClearedIsWaiting", + chat: database.Chat{ + WorkspaceID: workspaceID, + ContextAggregateHash: []byte{}, + }, + wantContext: true, + wantState: codersdk.ChatContextStateWaiting, + }, + { + name: "NonWorkspaceStaysNil", + chat: database.Chat{}, + wantContext: false, + }, + { + // A non-workspace chat with dirty remnants keeps its Context + // (existing behavior) but reports no lifecycle state. + name: "NonWorkspaceDirtyRemnantHasNoState", + chat: database.Chat{ + ContextDirtySince: sql.NullTime{Time: now, Valid: true}, + }, + wantContext: true, + wantState: "", + wantDirty: true, + }, + } + + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + tc.chat.ID = uuid.New() + tc.chat.OwnerID = uuid.New() + tc.chat.LastModelConfigID = uuid.New() + tc.chat.Status = database.ChatStatusWaiting + tc.chat.CreatedAt = now + tc.chat.UpdatedAt = now + + result := db2sdk.Chat(tc.chat, nil, nil) + if !tc.wantContext { + require.Nil(t, result.Context) + return + } + require.NotNil(t, result.Context) + require.Equal(t, tc.wantState, result.Context.State) + require.Equal(t, tc.wantDirty, result.Context.Dirty) + }) + } +} + func TestChat_NilFilesOmitted(t *testing.T) { t.Parallel() diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 10f6bcd855a..00e2837b572 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -5770,12 +5770,12 @@ func (q *querier) HasTemplateVersionsUsingCachedModuleFileInOrg(ctx context.Cont return q.db.HasTemplateVersionsUsingCachedModuleFileInOrg(ctx, arg) } -func (q *querier) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) error { +func (q *querier) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) ([]uuid.UUID, error) { // System-level operation: an agent context push fans hydration out // across every not-yet-pinned chat for the agent, so it authorizes at // the resource level rather than per-chat. if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChat); err != nil { - return err + return nil, err } return q.db.HydrateAgentChatsContext(ctx, arg) } diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index a2592c270ec..01ae2678625 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3897,12 +3897,12 @@ func (m queryMetricsStore) HasTemplateVersionsUsingCachedModuleFileInOrg(ctx con return r0, r1 } -func (m queryMetricsStore) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) error { +func (m queryMetricsStore) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) ([]uuid.UUID, error) { start := time.Now() - r0 := m.s.HydrateAgentChatsContext(ctx, arg) + r0, r1 := m.s.HydrateAgentChatsContext(ctx, arg) m.queryLatencies.WithLabelValues("HydrateAgentChatsContext").Observe(time.Since(start).Seconds()) m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "HydrateAgentChatsContext").Inc() - return r0 + return r0, r1 } func (m queryMetricsStore) IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) { diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index ba2884a0d44..9bec1ca6809 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -7289,11 +7289,12 @@ func (mr *MockStoreMockRecorder) HasTemplateVersionsUsingCachedModuleFileInOrg(c } // HydrateAgentChatsContext mocks base method. -func (m *MockStore) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) error { +func (m *MockStore) HydrateAgentChatsContext(ctx context.Context, arg database.HydrateAgentChatsContextParams) ([]uuid.UUID, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "HydrateAgentChatsContext", ctx, arg) - ret0, _ := ret[0].(error) - return ret0 + ret0, _ := ret[0].([]uuid.UUID) + ret1, _ := ret[1].(error) + return ret0, ret1 } // HydrateAgentChatsContext indicates an expected call of HydrateAgentChatsContext. diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 218688f7b40..fa13d110190 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -998,11 +998,13 @@ type sqlcQuerier interface { // a chat's pinned hash and pinned bodies are always written together. // Runs as a side effect of an agent push and of chat-create hydration, // so chats created before the agent was ready pick up the snapshot - // without a dirty event. The ON CONFLICT upsert is defensive: a - // not-yet-hydrated chat has no pinned rows, so it normally inserts. + // without a dirty event. Returns the hydrated chat IDs so the caller + // can emit context_ready watch events after the transaction commits. + // The ON CONFLICT upsert is defensive: a not-yet-hydrated chat has no + // pinned rows, so it normally inserts. // Does not bump chats.updated_at; the resource upsert's ON CONFLICT branch // sets chat_context_resources.updated_at on the rows it rewrites. - HydrateAgentChatsContext(ctx context.Context, arg HydrateAgentChatsContextParams) error + HydrateAgentChatsContext(ctx context.Context, arg HydrateAgentChatsContextParams) ([]uuid.UUID, error) // Increments generation_attempt and returns the resulting value. IncrementChatGenerationAttempt(ctx context.Context, id uuid.UUID) (int64, error) // Adds cost_micros to the spend for (user_id, effective_group_id, day). diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 008b6ef4dbb..1bac610bf84 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -1296,10 +1296,12 @@ func TestChatContextHydration(t *testing.T) { require.NoError(t, err) // Hydrate stamps only the NULL-hash chat for this agent. - require.NoError(t, db.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{ + hydrated, err := db.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{ AgentID: agent.ID, AggregateHash: hashH, - })) + }) + require.NoError(t, err) + require.Equal(t, []uuid.UUID{chatNull.ID}, hydrated, "hydrate returns only the chats it stamped") gotNull, err := db.GetChatByID(ctx, chatNull.ID) require.NoError(t, err) require.Equal(t, hashH, gotNull.ContextAggregateHash, "NULL-hash chat is hydrated") diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index cb2663638c7..17165c1a461 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -9514,41 +9514,44 @@ func (q *sqlQuerier) GetUserGroupSpendLimit(ctx context.Context, arg GetUserGrou return limit_micros, err } -const hydrateAgentChatsContext = `-- name: HydrateAgentChatsContext :exec +const hydrateAgentChatsContext = `-- name: HydrateAgentChatsContext :many WITH hydrated AS ( UPDATE chats SET - context_aggregate_hash = $2, - context_error = $3 - WHERE agent_id = $1::uuid + context_aggregate_hash = $1, + context_error = $2 + WHERE agent_id = $3::uuid AND archived = false AND context_aggregate_hash IS NULL RETURNING id +), +copied AS ( + INSERT INTO chat_context_resources ( + chat_id, source, body_kind, body, content_hash, size_bytes, status, error, source_path + ) + SELECT + hydrated.id, r.source, r.body_kind, r.body, r.content_hash, + r.size_bytes, r.status, r.error, r.source_path + FROM hydrated + CROSS JOIN workspace_agent_context_resources r + WHERE r.workspace_agent_id = $3::uuid + ON CONFLICT (chat_id, source) DO UPDATE SET + body_kind = EXCLUDED.body_kind, + body = EXCLUDED.body, + content_hash = EXCLUDED.content_hash, + size_bytes = EXCLUDED.size_bytes, + status = EXCLUDED.status, + error = EXCLUDED.error, + source_path = EXCLUDED.source_path, + updated_at = now() ) -INSERT INTO chat_context_resources ( - chat_id, source, body_kind, body, content_hash, size_bytes, status, error, source_path -) -SELECT - hydrated.id, r.source, r.body_kind, r.body, r.content_hash, - r.size_bytes, r.status, r.error, r.source_path -FROM hydrated -CROSS JOIN workspace_agent_context_resources r -WHERE r.workspace_agent_id = $1::uuid -ON CONFLICT (chat_id, source) DO UPDATE SET - body_kind = EXCLUDED.body_kind, - body = EXCLUDED.body, - content_hash = EXCLUDED.content_hash, - size_bytes = EXCLUDED.size_bytes, - status = EXCLUDED.status, - error = EXCLUDED.error, - source_path = EXCLUDED.source_path, - updated_at = now() +SELECT hydrated.id FROM hydrated ` type HydrateAgentChatsContextParams struct { - AgentID uuid.UUID `db:"agent_id" json:"agent_id"` AggregateHash []byte `db:"aggregate_hash" json:"aggregate_hash"` ContextError string `db:"context_error" json:"context_error"` + AgentID uuid.UUID `db:"agent_id" json:"agent_id"` } // Stamps the pinned hash and error on every not-yet-hydrated chat for @@ -9557,13 +9560,33 @@ type HydrateAgentChatsContextParams struct { // a chat's pinned hash and pinned bodies are always written together. // Runs as a side effect of an agent push and of chat-create hydration, // so chats created before the agent was ready pick up the snapshot -// without a dirty event. The ON CONFLICT upsert is defensive: a -// not-yet-hydrated chat has no pinned rows, so it normally inserts. +// without a dirty event. Returns the hydrated chat IDs so the caller +// can emit context_ready watch events after the transaction commits. +// The ON CONFLICT upsert is defensive: a not-yet-hydrated chat has no +// pinned rows, so it normally inserts. // Does not bump chats.updated_at; the resource upsert's ON CONFLICT branch // sets chat_context_resources.updated_at on the rows it rewrites. -func (q *sqlQuerier) HydrateAgentChatsContext(ctx context.Context, arg HydrateAgentChatsContextParams) error { - _, err := q.db.ExecContext(ctx, hydrateAgentChatsContext, arg.AgentID, arg.AggregateHash, arg.ContextError) - return err +func (q *sqlQuerier) HydrateAgentChatsContext(ctx context.Context, arg HydrateAgentChatsContextParams) ([]uuid.UUID, error) { + rows, err := q.db.QueryContext(ctx, hydrateAgentChatsContext, arg.AggregateHash, arg.ContextError, arg.AgentID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []uuid.UUID + for rows.Next() { + var id uuid.UUID + if err := rows.Scan(&id); err != nil { + return nil, err + } + items = append(items, id) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil } const incrementChatGenerationAttempt = `-- name: IncrementChatGenerationAttempt :one diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 97d2eeb9f7d..0b64780ecb3 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -1419,15 +1419,17 @@ SET context_dirty_since = NULL WHERE id = @id::uuid; --- name: HydrateAgentChatsContext :exec +-- name: HydrateAgentChatsContext :many -- Stamps the pinned hash and error on every not-yet-hydrated chat for -- an agent (context_aggregate_hash IS NULL) and copies the agent's -- current context resources onto those chats in the same statement, so -- a chat's pinned hash and pinned bodies are always written together. -- Runs as a side effect of an agent push and of chat-create hydration, -- so chats created before the agent was ready pick up the snapshot --- without a dirty event. The ON CONFLICT upsert is defensive: a --- not-yet-hydrated chat has no pinned rows, so it normally inserts. +-- without a dirty event. Returns the hydrated chat IDs so the caller +-- can emit context_ready watch events after the transaction commits. +-- The ON CONFLICT upsert is defensive: a not-yet-hydrated chat has no +-- pinned rows, so it normally inserts. -- Does not bump chats.updated_at; the resource upsert's ON CONFLICT branch -- sets chat_context_resources.updated_at on the rows it rewrites. WITH hydrated AS ( @@ -1439,25 +1441,28 @@ WITH hydrated AS ( AND archived = false AND context_aggregate_hash IS NULL RETURNING id +), +copied AS ( + INSERT INTO chat_context_resources ( + chat_id, source, body_kind, body, content_hash, size_bytes, status, error, source_path + ) + SELECT + hydrated.id, r.source, r.body_kind, r.body, r.content_hash, + r.size_bytes, r.status, r.error, r.source_path + FROM hydrated + CROSS JOIN workspace_agent_context_resources r + WHERE r.workspace_agent_id = @agent_id::uuid + ON CONFLICT (chat_id, source) DO UPDATE SET + body_kind = EXCLUDED.body_kind, + body = EXCLUDED.body, + content_hash = EXCLUDED.content_hash, + size_bytes = EXCLUDED.size_bytes, + status = EXCLUDED.status, + error = EXCLUDED.error, + source_path = EXCLUDED.source_path, + updated_at = now() ) -INSERT INTO chat_context_resources ( - chat_id, source, body_kind, body, content_hash, size_bytes, status, error, source_path -) -SELECT - hydrated.id, r.source, r.body_kind, r.body, r.content_hash, - r.size_bytes, r.status, r.error, r.source_path -FROM hydrated -CROSS JOIN workspace_agent_context_resources r -WHERE r.workspace_agent_id = @agent_id::uuid -ON CONFLICT (chat_id, source) DO UPDATE SET - body_kind = EXCLUDED.body_kind, - body = EXCLUDED.body, - content_hash = EXCLUDED.content_hash, - size_bytes = EXCLUDED.size_bytes, - status = EXCLUDED.status, - error = EXCLUDED.error, - source_path = EXCLUDED.source_path, - updated_at = now(); +SELECT hydrated.id FROM hydrated; -- name: MarkChatsContextDirtyByAgent :many -- Flips active, already-hydrated chats for an agent to dirty when the diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index d1519231ff6..2666f246480 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -825,6 +825,15 @@ The generation goroutine supports: - turn limit after a user message (the LLM shouldn't be able to spin forever in loop) - and other things +##### Context report gate + +A workspace-bound turn treats the agent's context report as a hard precondition. During generation preparation, before the pinned workspace context is read: + +- The chat's agent binding is resolved and, when the workspace has a newer start-transition build than the one the chat is bound to, rebound to the latest build's selected agent. Rebinding re-pins the chat's context (clearing it when the new agent has not pushed a snapshot yet). When the latest build is a stop or delete transition, the existing binding and pinned context are kept. +- A chat whose context is already pinned proceeds immediately; agent resolution errors are swallowed on that path so pinned chats on stopped (zero-agent) workspaces keep working from their pinned context. +- An unpinned chat waits on a 1-second poll loop, re-resolving the agent and checking for a pushed snapshot each tick. A push that arrives mid-wait is picked up on the next tick; the gate then pins the chat to it and the turn proceeds with that context. The wait is bounded by a 15-minute ceiling. +- The gate fails fast, without burning the ceiling, when the workspace's latest build is not a start transition or when the bound agent connected with an Agent API below 2.10 (the context-push minimum). Ceiling expiry and fast-fail conditions are terminal: the prepare phase does not retry them, and the message lands in the chat's persisted `last_error` as the visible error state. Context cancellation propagates unchanged so interrupts keep working during the wait. + ##### Reasoning effort Model configs may carry a `reasoning_effort` config (`{default, max}`) inside `chat_model_configs.options`. Users select a per-turn effort when sending or editing a message; the value is stored on `chat_messages.reasoning_effort` and on `chat_queued_messages.reasoning_effort` for queued messages. Queued messages carry the value through promotion, and `chats.last_reasoning_effort` tracks the most recent message that set one, mirroring `last_model_config_id`. diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index ccc313a89ee..ad6d3eaf757 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -189,6 +189,10 @@ type Server struct { streamSyncPoller *streamSyncPoller recordingSem chan struct{} + // contextReportTimeout is the ceiling on how long a turn waits for the + // workspace agent's first context report before failing the turn. + contextReportTimeout time.Duration + aibridgeTransportFactory *atomic.Pointer[aibridge.TransportFactory] experiments codersdk.Experiments @@ -650,17 +654,27 @@ func (c *turnWorkspaceContext) loadWorkspaceAgentLocked( if chatSnapshot.AgentID.Valid { agent, err := c.server.db.GetWorkspaceAgentByID(ctx, chatSnapshot.AgentID.UUID) if err == nil { - latestChat, workspaceMatches := c.currentWorkspaceMatches(chatSnapshot.WorkspaceID) - if !workspaceMatches { - chatSnapshot = latestChat - continue + stale, staleErr := c.chatBindingStale(ctx, chatSnapshot) + if staleErr != nil { + return chatSnapshot, database.WorkspaceAgent{}, staleErr } - c.agent = agent - c.agentLoaded = true - c.cachedWorkspaceID = chatSnapshot.WorkspaceID - return chatSnapshot, c.agent, nil - } - if !xerrors.Is(err, sql.ErrNoRows) { + if !stale { + latestChat, workspaceMatches := c.currentWorkspaceMatches(chatSnapshot.WorkspaceID) + if !workspaceMatches { + chatSnapshot = latestChat + continue + } + c.agent = agent + c.agentLoaded = true + c.cachedWorkspaceID = chatSnapshot.WorkspaceID + return chatSnapshot, c.agent, nil + } + // The workspace has a newer start build than the one the chat + // is bound to. GetWorkspaceAgentByID keeps resolving agents + // from old builds, so fall through to the re-resolve path + // below, which rebinds to the latest build's agent and + // re-pins the chat's context. + } else if !xerrors.Is(err, sql.ErrNoRows) { c.server.logger.Warn(ctx, "agent binding lookup failed, re-resolving", slog.F("agent_id", chatSnapshot.AgentID.UUID), slog.Error(err), @@ -721,6 +735,38 @@ func (c *turnWorkspaceContext) loadWorkspaceAgentLocked( ) } +// chatBindingStale reports whether the chat's bound agent belongs to an +// older build. Rebinding only chases start builds: when the latest build is +// a stop or delete transition, the chat keeps its existing agent and pinned +// context so pinned turns still work on a stopped workspace. +func (c *turnWorkspaceContext) chatBindingStale( + ctx context.Context, + chatSnapshot database.Chat, +) (bool, error) { + build, err := c.server.db.GetLatestWorkspaceBuildByWorkspaceID(ctx, chatSnapshot.WorkspaceID.UUID) + if err != nil { + return false, xerrors.Errorf("get latest workspace build: %w", err) + } + if build.Transition != database.WorkspaceTransitionStart { + return false, nil + } + return !chatSnapshot.BuildID.Valid || chatSnapshot.BuildID.UUID != build.ID, nil +} + +// refreshWorkspaceAgent drops the cached agent and re-resolves it, so the +// caller observes rebinds performed by loadWorkspaceAgentLocked while the +// cache would otherwise short-circuit. The cached workspace connection is +// left alone; the dial path revalidates it on the next use. +func (c *turnWorkspaceContext) refreshWorkspaceAgent( + ctx context.Context, +) (database.Chat, database.WorkspaceAgent, error) { + c.mu.Lock() + defer c.mu.Unlock() + c.agent = database.WorkspaceAgent{} + c.agentLoaded = false + return c.loadWorkspaceAgentLocked(ctx) +} + func (c *turnWorkspaceContext) latestWorkspaceAgentID( ctx context.Context, workspaceID uuid.UUID, @@ -3132,6 +3178,7 @@ func New(ps pubsub.Pubsub, cfg Config) *Server { agentConnFn: cfg.AgentConn, agentInactiveDisconnectTimeout: cfg.AgentInactiveDisconnectTimeout, dialTimeout: defaultDialTimeout, + contextReportTimeout: defaultContextReportTimeout, instructionLookupTimeout: instructionLookupTimeout, createWorkspaceFn: cfg.CreateWorkspace, startWorkspaceFn: cfg.StartWorkspace, diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 98adc25ccef..254f474f01a 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -1283,6 +1283,7 @@ func TestTurnWorkspaceContext_BindingFirstPath(t *testing.T) { workspaceAgent := database.WorkspaceAgent{ID: agentID} db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID).Return(workspaceAgent, nil).Times(1) + expectBoundBindingFresh(db, workspaceID) chatStateMu := &sync.Mutex{} currentChat := chat @@ -1372,6 +1373,16 @@ func expectBestEffortContextRepin(db *dbmock.MockStore) { db.EXPECT().DeleteChatContextResourcesByChatID(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() } +// expectBoundBindingFresh satisfies the bound-agent staleness probe that +// loadWorkspaceAgentLocked runs before trusting an existing agent binding. +// The latest build is reported without a start transition, so the probe +// keeps the chat's bound agent and the test exercises its original flow. +func expectBoundBindingFresh(db *dbmock.MockStore, workspaceID uuid.UUID) { + db.EXPECT().GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID). + Return(database.WorkspaceBuild{WorkspaceID: workspaceID}, nil). + AnyTimes() +} + func TestTurnWorkspaceContext_StaleBindingRepair(t *testing.T) { t.Parallel() @@ -1460,7 +1471,6 @@ func TestTurnWorkspaceContextGetWorkspaceConnLazyValidationSwitchesWorkspaceAgen gomock.InOrder( db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), staleAgentID).Return(staleAgent, nil), db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID).Return([]database.WorkspaceAgent{currentAgent}, nil), - db.EXPECT().GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID).Return(database.WorkspaceBuild{ID: buildID}, nil), db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), currentAgentID).Return(currentAgent, nil), db.EXPECT().UpdateChatBuildAgentBinding(gomock.Any(), database.UpdateChatBuildAgentBindingParams{ BuildID: uuid.NullUUID{UUID: buildID, Valid: true}, @@ -1468,6 +1478,10 @@ func TestTurnWorkspaceContextGetWorkspaceConnLazyValidationSwitchesWorkspaceAgen ID: chat.ID, }).Return(updatedChat, nil), ) + // Serves both the bound-agent staleness probe (the zero transition is + // not a start, so the binding is kept) and the post-switch build fetch. + db.EXPECT().GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID). + Return(database.WorkspaceBuild{ID: buildID}, nil).AnyTimes() conn := agentconnmock.NewMockAgentConn(ctrl) conn.EXPECT().SetExtraHeaders(gomock.Any()).Times(1) @@ -1535,6 +1549,7 @@ func TestTurnWorkspaceContextGetWorkspaceConnFastFailsWithoutCurrentAgent(t *tes db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), staleAgentID). Return(staleAgent, nil). Times(1) + expectBoundBindingFresh(db, workspaceID) db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{}, nil). Times(1) @@ -2223,9 +2238,10 @@ func TestGetWorkspaceConn_StaleAgentRecovery(t *testing.T) { // Lazy validation discovers the new agent. db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{newAgent}, nil).Times(1) - // Post-switch: persist the new binding. + // Serves both the bound-agent staleness probe (the zero transition is + // not a start, so the binding is kept) and the post-switch build fetch. db.EXPECT().GetLatestWorkspaceBuildByWorkspaceID(gomock.Any(), workspaceID). - Return(database.WorkspaceBuild{ID: buildID}, nil).Times(1) + Return(database.WorkspaceBuild{ID: buildID}, nil).AnyTimes() db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), newAgentID). Return(newAgent, nil).Times(1) @@ -2331,6 +2347,7 @@ func TestGetWorkspaceConn_SameBuildAgentCrash(t *testing.T) { // ensureWorkspaceAgent fetches the (crashed) agent. db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). Return(agent, nil).Times(1) + expectBoundBindingFresh(db, workspaceID) // Validation finds the same agent in the latest build. db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{agent}, nil).Times(1) @@ -2598,6 +2615,7 @@ func TestGetWorkspaceConn_DialTimeoutDisconnectedRecoveryThreshold(t *testing.T) db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). Return(disconnectedAgent, nil). Times(2) + expectBoundBindingFresh(db, workspaceID) db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{disconnectedAgent}, nil). Times(1) @@ -2710,6 +2728,7 @@ func TestGetWorkspaceConn_DisconnectedStatusDialSuccessDoesNotEscalate(t *testin db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). Return(disconnectedAgent, nil). Times(1) + expectBoundBindingFresh(db, workspaceID) server := &Server{ db: db, @@ -2779,6 +2798,7 @@ func TestGetWorkspaceConn_CacheHitDisconnectedRetriesDialBeforeEscalating(t *tes db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). Return(disconnectedAgent, nil). Times(2) + expectBoundBindingFresh(db, workspaceID) server := &Server{ db: db, @@ -2859,6 +2879,7 @@ func TestGetWorkspaceConn_DialTimeout(t *testing.T) { db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). Return(connectedAgent, nil). Times(2) + expectBoundBindingFresh(db, workspaceID) db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{connectedAgent}, nil). Times(1) @@ -2923,6 +2944,7 @@ func TestGetWorkspaceConn_DialTimeoutStatusTimeoutDoesNotEscalate(t *testing.T) db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). Return(timedOutAgent, nil). Times(2) + expectBoundBindingFresh(db, workspaceID) db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{timedOutAgent}, nil). Times(1) @@ -2994,6 +3016,7 @@ func TestGetWorkspaceConn_DialTimeoutParentCanceled(t *testing.T) { db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). Return(connectedAgent, nil). Times(1) + expectBoundBindingFresh(db, workspaceID) parentErr := xerrors.New("parent canceled") ctx, cancel := context.WithCancelCause(testutil.Context(t, testutil.WaitShort)) @@ -3074,6 +3097,7 @@ func TestGetWorkspaceConn_PreflightExternalAgentTimedOut(t *testing.T) { db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). Return(agent, nil). Times(1) + expectBoundBindingFresh(db, workspaceID) db.EXPECT().GetWorkspaceAgentsInLatestBuildByWorkspaceID(gomock.Any(), workspaceID). Return([]database.WorkspaceAgent{agent}, nil). Times(1) @@ -3148,6 +3172,7 @@ func TestGetWorkspaceConn_PreflightExternalAgentConnectingDials(t *testing.T) { db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). Return(agent, nil). Times(1) + expectBoundBindingFresh(db, workspaceID) conn := agentconnmock.NewMockAgentConn(ctrl) conn.EXPECT().SetExtraHeaders(gomock.Any()).Times(1) @@ -3224,6 +3249,7 @@ func TestGetWorkspaceConn_DialErrorNotMisclassifiedAsTimeout(t *testing.T) { db.EXPECT().GetWorkspaceAgentByID(gomock.Any(), agentID). Return(connectedAgent, nil). Times(1) + expectBoundBindingFresh(db, workspaceID) // When the initial dial fails immediately, dialWithLazyValidation // calls resolveFastFailure which validates the binding. Mock the // validation to return the same agent, triggering a synchronous diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 300ed4c647d..9c24351073f 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -45,6 +45,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbfake" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/database/dbtime" dbpubsub "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/coderd/util/ptr" "github.com/coder/coder/v2/coderd/util/slice" @@ -4862,7 +4863,15 @@ func TestCreateWorkspaceTool_EndToEnd(t *testing.T) { require.True(t, foundToolResultInSecondCall, "expected second streamed model call to include create_workspace tool output") } -func TestStartWorkspaceTool_EndToEnd(t *testing.T) { +// TestStoppedWorkspaceUnpinnedChatSurfacesContextGateError covers the +// context-report gate end to end through the chat worker: a chat bound to a +// stopped workspace that never reported context cannot run a turn, and the +// gate's actionable error surfaces as the chat's visible error state instead +// of an opaque retry loop. Before the gate, this flow ran the turn and let +// the model call start_workspace; the started-workspace requirement now +// front-runs the model call, and the start_workspace tool behavior itself is +// covered by chattool's tests. +func TestStoppedWorkspaceUnpinnedChatSurfacesContextGateError(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitSuperLong) @@ -4882,11 +4891,8 @@ func TestStartWorkspaceTool_EndToEnd(t *testing.T) { coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID) - // Create a workspace, then stop it so start_workspace has - // something to start. We intentionally skip starting a test - // agent. The echo provisioner creates new agent rows for each - // build, so an agent started for build 1 cannot serve build 3. - // The tool handles the no-agent case gracefully. + // Create a workspace, then stop it. No agent ever connects or pushes + // context, so a chat bound to it has no pinned context to run from. workspace := coderdtest.CreateWorkspace(t, client, template.ID) coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) workspace = coderdtest.MustTransitionWorkspace( @@ -4894,26 +4900,12 @@ func TestStartWorkspaceTool_EndToEnd(t *testing.T) { codersdk.WorkspaceTransitionStart, codersdk.WorkspaceTransitionStop, ) - var streamedCallCount atomic.Int32 - var streamedCallsMu sync.Mutex - streamedCalls := make([][]chattest.OpenAIMessage, 0, 2) - openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { if !req.Stream { return chattest.OpenAINonStreamingResponse("Start workspace test") } - - streamedCallsMu.Lock() - streamedCalls = append(streamedCalls, append([]chattest.OpenAIMessage(nil), req.Messages...)) - streamedCallsMu.Unlock() - - if streamedCallCount.Add(1) == 1 { - return chattest.OpenAIStreamingResponse( - chattest.OpenAIToolCallChunk("start_workspace", "{}"), - ) - } return chattest.OpenAIStreamingResponse( - chattest.OpenAITextChunks("Workspace started and ready.")..., + chattest.OpenAITextChunks("This turn must not run.")..., ) }) @@ -4942,69 +4934,11 @@ func TestStartWorkspaceTool_EndToEnd(t *testing.T) { return got.Status == codersdk.ChatStatusWaiting || got.Status == codersdk.ChatStatusError }, testutil.WaitSuperLong, testutil.IntervalFast) - if chatResult.Status == codersdk.ChatStatusError { - lastError := "" - if chatResult.LastError != nil { - lastError = chatResult.LastError.Message - } - require.FailNowf(t, "chat run failed", "last_error=%q", lastError) - } - - // Verify the workspace was started. - require.NotNil(t, chatResult.WorkspaceID) - updatedWorkspace, err := client.Workspace(ctx, workspace.ID) - require.NoError(t, err) - require.Equal(t, codersdk.WorkspaceTransitionStart, updatedWorkspace.LatestBuild.Transition) - - chatMsgs, err := expClient.GetChatMessages(ctx, chat.ID, nil) - require.NoError(t, err) - - // Verify start_workspace tool result exists in the chat messages. - var foundStartWorkspaceResult bool - for _, message := range chatMsgs.Messages { - if message.Role != codersdk.ChatMessageRoleTool { - continue - } - for _, part := range message.Content { - if part.Type != codersdk.ChatMessagePartTypeToolResult || part.ToolName != "start_workspace" { - continue - } - var result map[string]any - require.NoError(t, json.Unmarshal(part.Result, &result)) - started, ok := result["started"].(bool) - require.True(t, ok) - require.True(t, started) - foundStartWorkspaceResult = true - } - } - require.True(t, foundStartWorkspaceResult, "expected start_workspace tool result message") - - // Verify the LLM received the tool result in its second call. - require.GreaterOrEqual(t, streamedCallCount.Load(), int32(2)) - streamedCallsMu.Lock() - recordedStreamCalls := append([][]chattest.OpenAIMessage(nil), streamedCalls...) - streamedCallsMu.Unlock() - require.GreaterOrEqual(t, len(recordedStreamCalls), 2) - - var foundToolResultInSecondCall bool - for _, message := range recordedStreamCalls[1] { - if message.Role != "tool" { - continue - } - if !json.Valid([]byte(message.Content)) { - continue - } - var result map[string]any - if err := json.Unmarshal([]byte(message.Content), &result); err != nil { - continue - } - started, ok := result["started"].(bool) - if ok && started { - foundToolResultInSecondCall = true - break - } - } - require.True(t, foundToolResultInSecondCall, "expected second streamed model call to include start_workspace tool output") + require.Equal(t, codersdk.ChatStatusError, chatResult.Status, + "an unpinned chat on a stopped workspace must fail the turn visibly") + require.NotNil(t, chatResult.LastError) + require.Contains(t, chatResult.LastError.Message, + "workspace must be started to report chat context") } func TestStoppedWorkspaceWithPersistedAgentBindingDoesNotBlockChat(t *testing.T) { @@ -5074,6 +5008,17 @@ func TestStoppedWorkspaceWithPersistedAgentBindingDoesNotBlockChat(t *testing.T) }) require.NoError(t, err) + // The agent reported context while the workspace ran, pinning the chat. + // The context-report gate only lets a pinned chat run once the workspace + // is stopped below, matching the production sequence this test models. + snapshot, err := db.GetLatestWorkspaceAgentContextSnapshot(ctx, dbAgent.ID) + require.NoError(t, err) + _, err = db.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{ + AgentID: dbAgent.ID, + AggregateHash: snapshot.AggregateHash, + }) + require.NoError(t, err) + dbfake.WorkspaceBuild(t, db, ws).Seed(database.WorkspaceBuild{ Transition: database.WorkspaceTransitionStop, BuildNumber: 2, @@ -8696,7 +8641,18 @@ func seedWorkspaceWithAgent( Version: "v1.0.0", ExpandedDirectory: "/home/coder/project", })) - dbAgent, err := db.GetWorkspaceAgentByID(context.Background(), dbAgent.ID) + // The context-report gate blocks workspace turns until the agent has + // pushed a context snapshot. Seed an empty snapshot so turns proceed + // as they did before the gate; row existence alone marks the agent as + // having reported, and an empty snapshot contributes no prompt context. + _, err := db.UpsertWorkspaceAgentContextSnapshot(context.Background(), database.UpsertWorkspaceAgentContextSnapshotParams{ + WorkspaceAgentID: dbAgent.ID, + Version: 1, + AggregateHash: []byte{0x01}, + ReceivedAt: dbtime.Now(), + }) + require.NoError(t, err) + dbAgent, err = db.GetWorkspaceAgentByID(context.Background(), dbAgent.ID) require.NoError(t, err) return ws, dbAgent } diff --git a/coderd/x/chatd/context_gate.go b/coderd/x/chatd/context_gate.go new file mode 100644 index 00000000000..4a24b56d488 --- /dev/null +++ b/coderd/x/chatd/context_gate.go @@ -0,0 +1,296 @@ +package chatd + +import ( + "context" + "time" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "cdr.dev/slog/v3" + "github.com/coder/coder/v2/apiversion" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/x/chatd/chaterror" + "github.com/coder/coder/v2/codersdk" +) + +const ( + // contextReportPollInterval is how often the context-report gate + // re-resolves the chat's agent and re-checks for a pushed snapshot + // while a turn waits for the agent's first context report. + contextReportPollInterval = time.Second + // defaultContextReportTimeout bounds how long a turn waits for the + // workspace agent to report its context snapshot before failing the + // turn with a visible error. + defaultContextReportTimeout = 15 * time.Minute + + contextReportTickerTag = "context-report-tick" + contextReportTimeoutTag = "context-report-timeout" + contextReportTimerTagGroup = "chatd" + + // Context pushes require Agent API >= 2.10 (PushContextState). Agents + // built before that version can never report context, so waiting on + // them is pointless. + contextReportMinAPIMajor = 2 + contextReportMinAPIMinor = 10 +) + +var ( + errChatContextWorkspaceNotStarted = newChatContextGateError( + "workspace must be started to report chat context") + errChatContextAgentTooOld = newChatContextGateError( + "workspace agent is too old to report chat context; " + + "update Coder and rebuild the workspace") +) + +// newChatContextGateError builds a context-gate error whose text survives +// chaterror classification, so the exact message lands in the chat's +// persisted last_error instead of a generic fallback. +func newChatContextGateError(message string) error { + return chaterror.WithClassification( + xerrors.New(message), + chaterror.ClassifiedError{ + Message: message, + Kind: codersdk.ChatErrorKindGeneric, + }, + ) +} + +// awaitChatContextReported enforces that a workspace-bound turn does not run +// before the chat's agent has reported a context snapshot. A chat whose +// current snapshot is already pinned proceeds immediately; agent resolution +// errors are swallowed on that path so pinned chats on stopped (zero-agent) +// workspaces keep working. An unpinned chat waits on a poll loop that +// re-resolves the agent (rebinding to a newer start build when the binding is +// stale) and exits as soon as the bound agent has any snapshot row, then pins +// the chat to it. +// +// The wait fails fast, without burning the ceiling, when the workspace's +// latest build is not a start transition or when the bound agent connected +// with an Agent API too old to push context. Context cancellation propagates +// unchanged so interrupts keep working. Terminal failures are wrapped with +// terminalGeneration so the prepare phase does not retry them and the message +// surfaces as the chat's visible error state. +func (server *Server) awaitChatContextReported( + ctx context.Context, + workspaceCtx *turnWorkspaceContext, + logger slog.Logger, +) (database.WorkspaceAgent, error) { + before := workspaceCtx.currentChatSnapshot() + resolvedChat, agent, resolveErr := workspaceCtx.ensureWorkspaceAgent(ctx) + if resolveErr == nil && + (!nullUUIDEqual(before.AgentID, resolvedChat.AgentID) || + !nullUUIDEqual(before.BuildID, resolvedChat.BuildID)) { + // The resolution rebound the chat to another build's agent and + // re-pinned its context in a separate transaction, so the cached + // snapshot's pinned fields are stale. Reload before reading them. + fresh, err := workspaceCtx.loadChatSnapshot(ctx, resolvedChat.ID) + if err != nil { + return database.WorkspaceAgent{}, xerrors.Errorf( + "reload chat after agent rebind: %w", err) + } + workspaceCtx.setCurrentChat(fresh) + } + + if chatContextPinned(workspaceCtx.currentChatSnapshot()) { + // Pinned and, because ensureWorkspaceAgent rebinds stale bindings + // and rebinding re-pins, current. Proceed as before the gate + // existed: resolveErr is swallowed so a pinned chat on a stopped + // workspace with no agents still runs from its pinned context. + if resolveErr != nil { + logger.Debug(ctx, "context gate: proceeding with pinned context despite agent resolution error", + slog.Error(resolveErr)) + } + return agent, nil + } + if ctx.Err() != nil { + return database.WorkspaceAgent{}, ctx.Err() + } + + agent, done, err := server.checkChatContextReported(ctx, workspaceCtx, logger) + if err != nil { + return database.WorkspaceAgent{}, err + } + if done { + return agent, nil + } + + logger.Info(ctx, "context gate: waiting for workspace agent to report chat context", + slog.F("timeout", server.contextReportTimeout)) + + timeout := server.clock.NewTimer(server.contextReportTimeout, + contextReportTimerTagGroup, contextReportTimeoutTag) + defer timeout.Stop() + ticker := server.clock.NewTicker(contextReportPollInterval, + contextReportTimerTagGroup, contextReportTickerTag) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return database.WorkspaceAgent{}, ctx.Err() + case <-timeout.C: + return database.WorkspaceAgent{}, terminalGeneration(newChatContextGateError( + "workspace agent did not report chat context within " + + server.contextReportTimeout.String())) + case <-ticker.C: + agent, done, err := server.checkChatContextReported(ctx, workspaceCtx, logger) + if err != nil { + return database.WorkspaceAgent{}, err + } + if done { + return agent, nil + } + } + } +} + +// checkChatContextReported performs one context-gate poll: it re-resolves +// (and possibly rebinds) the chat's agent, applies the fast-fail conditions, +// and reports done=true once the bound agent has a pushed snapshot, at which +// point the chat is pinned to it. done=false with a nil error means "keep +// waiting". +func (server *Server) checkChatContextReported( + ctx context.Context, + workspaceCtx *turnWorkspaceContext, + logger slog.Logger, +) (database.WorkspaceAgent, bool, error) { + chatSnapshot, agent, err := workspaceCtx.refreshWorkspaceAgent(ctx) + if err != nil { + if ctx.Err() != nil { + return database.WorkspaceAgent{}, false, ctx.Err() + } + if !xerrors.Is(err, errChatHasNoWorkspaceAgent) { + return database.WorkspaceAgent{}, false, err + } + // The latest build has no agent rows. A start build that is still + // provisioning will grow them, so keep waiting; any other + // transition never will, so fail fast. + started, buildErr := server.latestBuildIsStart(ctx, chatSnapshot.WorkspaceID) + if buildErr != nil { + return database.WorkspaceAgent{}, false, buildErr + } + if !started { + return database.WorkspaceAgent{}, false, terminalGeneration(errChatContextWorkspaceNotStarted) + } + return database.WorkspaceAgent{}, false, nil + } + + // A resolved binding is kept as-is when the latest build is not a + // start transition, so an unpinned chat on a stopped workspace would + // otherwise wait the full ceiling for a push that cannot arrive. + started, buildErr := server.latestBuildIsStart(ctx, chatSnapshot.WorkspaceID) + if buildErr != nil { + return database.WorkspaceAgent{}, false, buildErr + } + if !started { + return database.WorkspaceAgent{}, false, terminalGeneration(errChatContextWorkspaceNotStarted) + } + + if agentTooOldForContextReport(ctx, agent, logger) { + return database.WorkspaceAgent{}, false, terminalGeneration(errChatContextAgentTooOld) + } + + //nolint:gocritic // Chatd reads the agent's snapshot as the daemon subject. + _, _, hasSnapshot, err := latestAgentSnapshot(dbauthz.AsChatd(ctx), server.db, agent.ID) + if err != nil { + return database.WorkspaceAgent{}, false, err + } + if !hasSnapshot { + return database.WorkspaceAgent{}, false, nil + } + + // The snapshot may have hydrated the chat mid-wait (pushes stamp bound + // NULL-hash chats), so reload the row before pinning. + fresh, err := workspaceCtx.loadChatSnapshot(ctx, chatSnapshot.ID) + if err != nil { + return database.WorkspaceAgent{}, false, xerrors.Errorf( + "reload chat after context report: %w", err) + } + if !chatContextPinned(fresh) { + if fresh.ContextAggregateHash == nil { + // A never-pinned chat carries a NULL hash, which the push-side + // NULL-gated hydrate can stamp; this also never clobbers a + // concurrent push that hydrated the chat first. + server.ensureChatContextPinnedOnFirstTurn(ctx, fresh) + } else { + // A rebind-cleared pin is an empty, non-NULL hash (the postgres + // driver encodes the nil clear as an empty bytea), which the + // NULL-gated hydrate skips. Re-pin explicitly, mirroring the + // refresh endpoint. + //nolint:gocritic // Chatd re-pins chats it does not own as the daemon subject. + repinCtx := dbauthz.AsChatd(ctx) + if err := database.ReadModifyUpdate(server.db, func(tx database.Store) error { + return repinChatContext(repinCtx, tx, fresh.ID, fresh.AgentID) + }); err != nil { + return database.WorkspaceAgent{}, false, xerrors.Errorf( + "pin chat to reported context: %w", err) + } + } + fresh, err = workspaceCtx.loadChatSnapshot(ctx, chatSnapshot.ID) + if err != nil { + return database.WorkspaceAgent{}, false, xerrors.Errorf( + "reload chat after context pin: %w", err) + } + // The pin committed here (not on the push path, whose own hydrate + // publishes for the chats it stamps), so announce the freshly + // pinned context to watchers. + if chatContextPinned(fresh) { + server.publishChatPubsubEvents( + []database.Chat{fresh}, codersdk.ChatWatchEventKindContextReady) + } + } + workspaceCtx.setCurrentChat(fresh) + return agent, true, nil +} + +// chatContextPinned reports whether the chat carries a meaningful pinned +// context hash. A NULL hash (never hydrated) and an empty non-NULL hash (a +// rebind cleared the pin; the postgres driver stores the nil clear as an +// empty bytea) both count as unpinned. +func chatContextPinned(chat database.Chat) bool { + return len(chat.ContextAggregateHash) > 0 +} + +// latestBuildIsStart reports whether the workspace's most recent build is a +// start transition. +func (server *Server) latestBuildIsStart( + ctx context.Context, + workspaceID uuid.NullUUID, +) (bool, error) { + if !workspaceID.Valid { + return false, xerrors.New("chat has no workspace") + } + build, err := server.db.GetLatestWorkspaceBuildByWorkspaceID(ctx, workspaceID.UUID) + if err != nil { + return false, xerrors.Errorf("get latest workspace build: %w", err) + } + return build.Transition == database.WorkspaceTransitionStart, nil +} + +// agentTooOldForContextReport reports whether the agent connected with an +// Agent API version below the context-push minimum. api_version is empty +// until the agent's first connect, so an agent that has never connected is +// not considered too old; the gate keeps waiting for it instead. An +// unparsable version is logged and treated the same way. +func agentTooOldForContextReport( + ctx context.Context, + agent database.WorkspaceAgent, + logger slog.Logger, +) bool { + if agent.APIVersion == "" { + return false + } + major, minor, err := apiversion.Parse(agent.APIVersion) + if err != nil { + logger.Warn(ctx, "context gate: unparsable agent api version", + slog.F("agent_id", agent.ID), + slog.F("api_version", agent.APIVersion), + slog.Error(err)) + return false + } + if major != contextReportMinAPIMajor { + return major < contextReportMinAPIMajor + } + return minor < contextReportMinAPIMinor +} diff --git a/coderd/x/chatd/context_gate_internal_test.go b/coderd/x/chatd/context_gate_internal_test.go new file mode 100644 index 00000000000..fec055d0c64 --- /dev/null +++ b/coderd/x/chatd/context_gate_internal_test.go @@ -0,0 +1,640 @@ +package chatd + +import ( + "context" + "database/sql" + "encoding/json" + "sync" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "google.golang.org/protobuf/encoding/protojson" + "google.golang.org/protobuf/types/known/structpb" + + "cdr.dev/slog/v3/sloggers/slogtest" + agentproto "github.com/coder/coder/v2/agent/proto" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/database/dbtime" + coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" + "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" + "github.com/coder/coder/v2/coderd/x/chatd/chatstate" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/testutil" + "github.com/coder/quartz" +) + +// newGateServer builds a minimal Server for exercising the context-report +// gate directly against a real database, with a caller-provided clock and +// wait ceiling. The pubsub carries the context_ready watch events the gate +// publishes at gate-exit pinning. +func newGateServer(t *testing.T, fix rebindFixture, clk quartz.Clock, ceiling time.Duration) *Server { + t.Helper() + return &Server{ + db: fix.db, + pubsub: fix.ps, + logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}), + clock: clk, + contextReportTimeout: ceiling, + } +} + +// subscribeChatWatchEvents collects the owner's chat watch events published +// via pubsub so tests can assert on gate-exit context_ready publishes. +func subscribeChatWatchEvents(t *testing.T, fix rebindFixture) <-chan codersdk.ChatWatchEvent { + t.Helper() + events := make(chan codersdk.ChatWatchEvent, 16) + cancel, err := fix.ps.Subscribe( + coderdpubsub.ChatWatchEventChannel(fix.user.ID), + func(_ context.Context, message []byte) { + var event codersdk.ChatWatchEvent + if err := json.Unmarshal(message, &event); err != nil { + return + } + events <- event + }) + require.NoError(t, err) + t.Cleanup(cancel) + return events +} + +func newGateTurnContext(t *testing.T, server *Server, chat database.Chat) *turnWorkspaceContext { + t.Helper() + cur := chat + wc := &turnWorkspaceContext{ + server: server, + chatStateMu: &sync.Mutex{}, + currentChat: &cur, + loadChatSnapshot: server.db.GetChatByID, + } + t.Cleanup(wc.close) + return wc +} + +type gateAwaitResult struct { + agent database.WorkspaceAgent + err error +} + +// awaitGateAsync runs awaitChatContextReported in a goroutine so the test can +// drive the mock clock while the gate waits. +func awaitGateAsync(ctx context.Context, server *Server, wc *turnWorkspaceContext) <-chan gateAwaitResult { + results := make(chan gateAwaitResult, 1) + go func() { + agent, err := server.awaitChatContextReported(ctx, wc, server.logger) + results <- gateAwaitResult{agent: agent, err: err} + }() + return results +} + +// seedStopBuild appends a stop-transition build so the workspace's latest +// build is no longer a start. +func seedStopBuild(t *testing.T, db database.Store, fix rebindFixture) { + t.Helper() + tv, err := db.GetTemplateVersionByID(fix.ctx, mustLatestBuild(t, db, fix).TemplateVersionID) + require.NoError(t, err) + job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + OrganizationID: fix.org.ID, + CompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, + }) + dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: fix.ws.ID, + TemplateVersionID: tv.ID, + JobID: job.ID, + BuildNumber: mustLatestBuild(t, db, fix).BuildNumber + 1, + Transition: database.WorkspaceTransitionStop, + }) +} + +// seedStartBuildWithAgent appends a start-transition build carrying a single +// fresh agent, simulating a workspace rebuild. +func seedStartBuildWithAgent(t *testing.T, db database.Store, fix rebindFixture) (database.WorkspaceBuild, database.WorkspaceAgent) { + t.Helper() + tv, err := db.GetTemplateVersionByID(fix.ctx, mustLatestBuild(t, db, fix).TemplateVersionID) + require.NoError(t, err) + job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + OrganizationID: fix.org.ID, + CompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, + }) + build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: fix.ws.ID, + TemplateVersionID: tv.ID, + JobID: job.ID, + BuildNumber: mustLatestBuild(t, db, fix).BuildNumber + 1, + Transition: database.WorkspaceTransitionStart, + }) + res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + Transition: database.WorkspaceTransitionStart, + JobID: job.ID, + }) + agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: res.ID}) + return build, agent +} + +func mustLatestBuild(t *testing.T, db database.Store, fix rebindFixture) database.WorkspaceBuild { + t.Helper() + build, err := db.GetLatestWorkspaceBuildByWorkspaceID(fix.ctx, fix.ws.ID) + require.NoError(t, err) + return build +} + +func gateChat(t *testing.T, fix rebindFixture, agentID uuid.UUID, buildID uuid.UUID) database.Chat { + t.Helper() + seed := database.Chat{ + OwnerID: fix.user.ID, + OrganizationID: fix.org.ID, + LastModelConfigID: fix.model.ID, + WorkspaceID: uuid.NullUUID{UUID: fix.ws.ID, Valid: true}, + Status: database.ChatStatusRunning, + } + if agentID != uuid.Nil { + seed.AgentID = uuid.NullUUID{UUID: agentID, Valid: true} + } + if buildID != uuid.Nil { + seed.BuildID = uuid.NullUUID{UUID: buildID, Valid: true} + } + return dbgen.Chat(t, fix.db, seed) +} + +func TestAwaitChatContextReported(t *testing.T) { + t.Parallel() + + // SnapshotAlreadyExists: the bound agent already pushed a snapshot, so + // the gate exits on its pre-wait check without any ticker interaction + // (the mock clock is never advanced), pins the chat, and announces the + // pin with a context_ready watch event. + t.Run("SnapshotAlreadyExists", func(t *testing.T) { + t.Parallel() + fix := newRebindFixture(t) + server := newGateServer(t, fix, quartz.NewMock(t), defaultContextReportTimeout) + chat := gateChat(t, fix, fix.agentA, fix.buildID) + require.Nil(t, chat.ContextAggregateHash) + events := subscribeChatWatchEvents(t, fix) + + wc := newGateTurnContext(t, server, chat) + agent, err := server.awaitChatContextReported(fix.ctx, wc, server.logger) + require.NoError(t, err) + require.Equal(t, fix.agentA, agent.ID) + + post, err := fix.db.GetChatByID(fix.ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, fix.hashA, post.ContextAggregateHash, "gate exit pins the chat to the reported snapshot") + + event := testutil.RequireReceive(fix.ctx, t, events) + require.Equal(t, codersdk.ChatWatchEventKindContextReady, event.Kind, + "gate-exit pinning publishes a context_ready event") + require.Equal(t, chat.ID, event.Chat.ID) + require.NotNil(t, event.Chat.Context) + require.Equal(t, codersdk.ChatContextStateReady, event.Chat.Context.State) + }) + + // PinnedSameBuild: a pinned chat whose binding matches the latest start + // build skips the gate entirely, even though its agent has no snapshot + // row (an entered gate could never exit here). + t.Run("PinnedSameBuild", func(t *testing.T) { + t.Parallel() + fix := newRebindFixture(t) + server := newGateServer(t, fix, quartz.NewMock(t), defaultContextReportTimeout) + chat := gateChat(t, fix, fix.agentNoSnap, fix.buildID) + require.NoError(t, fix.db.SetChatContextSnapshot(fix.ctx, database.SetChatContextSnapshotParams{ + ID: chat.ID, + AggregateHash: []byte{0xfe}, + })) + chat, err := fix.db.GetChatByID(fix.ctx, chat.ID) + require.NoError(t, err) + require.NotNil(t, chat.ContextAggregateHash) + + wc := newGateTurnContext(t, server, chat) + agent, err := server.awaitChatContextReported(fix.ctx, wc, server.logger) + require.NoError(t, err) + require.Equal(t, fix.agentNoSnap, agent.ID) + }) + + // PinnedStopBuild: a pinned chat on a workspace whose latest build is a + // stop transition keeps its binding and proceeds from pinned context, as + // before the gate existed. + t.Run("PinnedStopBuild", func(t *testing.T) { + t.Parallel() + fix := newRebindFixture(t) + server := newGateServer(t, fix, quartz.NewMock(t), defaultContextReportTimeout) + chat := gateChat(t, fix, fix.agentA, fix.buildID) + _, hydrateErr := fix.db.HydrateAgentChatsContext(fix.ctx, database.HydrateAgentChatsContextParams{ + AgentID: fix.agentA, + AggregateHash: fix.hashA, + }) + require.NoError(t, hydrateErr) + seedStopBuild(t, fix.db, fix) + + chat, err := fix.db.GetChatByID(fix.ctx, chat.ID) + require.NoError(t, err) + wc := newGateTurnContext(t, server, chat) + agent, err := server.awaitChatContextReported(fix.ctx, wc, server.logger) + require.NoError(t, err) + require.Equal(t, fix.agentA, agent.ID, "the stop build does not steal the binding") + + post, err := fix.db.GetChatByID(fix.ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, fix.hashA, post.ContextAggregateHash, "pinned context survives the stop build") + }) + + // StopBuildUnpinned: an unpinned chat cannot receive a context report + // from a workspace whose latest build is not a start, so the gate fails + // immediately with the actionable started-workspace error, both when the + // old agent binding still resolves and when no agent is bound at all. + t.Run("StopBuildUnpinned", func(t *testing.T) { + t.Parallel() + for _, bound := range []bool{true, false} { + name := "Bound" + if !bound { + name = "Unbound" + } + t.Run(name, func(t *testing.T) { + t.Parallel() + fix := newRebindFixture(t) + server := newGateServer(t, fix, quartz.NewMock(t), defaultContextReportTimeout) + agentID := uuid.Nil + buildID := uuid.Nil + if bound { + agentID = fix.agentNoSnap + buildID = fix.buildID + } + chat := gateChat(t, fix, agentID, buildID) + seedStopBuild(t, fix.db, fix) + + wc := newGateTurnContext(t, server, chat) + _, err := server.awaitChatContextReported(fix.ctx, wc, server.logger) + require.ErrorIs(t, err, errChatContextWorkspaceNotStarted) + require.ErrorIs(t, err, errTerminalGeneration, "the prepare phase must not retry the gate failure") + require.ErrorContains(t, err, "workspace must be started to report chat context") + }) + } + }) + + // AgentTooOld: an agent that connected with an Agent API below 2.10 can + // never push context, so the gate fails immediately instead of waiting + // out the ceiling. + t.Run("AgentTooOld", func(t *testing.T) { + t.Parallel() + fix := newRebindFixture(t) + server := newGateServer(t, fix, quartz.NewMock(t), defaultContextReportTimeout) + require.NoError(t, fix.db.UpdateWorkspaceAgentStartupByID(fix.ctx, database.UpdateWorkspaceAgentStartupByIDParams{ + ID: fix.agentNoSnap, + Version: "v2.0.0", + APIVersion: "2.0", + Subsystems: []database.WorkspaceAgentSubsystem{}, + })) + chat := gateChat(t, fix, fix.agentNoSnap, fix.buildID) + + wc := newGateTurnContext(t, server, chat) + _, err := server.awaitChatContextReported(fix.ctx, wc, server.logger) + require.ErrorIs(t, err, errChatContextAgentTooOld) + require.ErrorIs(t, err, errTerminalGeneration) + require.ErrorContains(t, err, "workspace agent is too old to report chat context") + }) + + // WaitsForPush: an unpinned chat whose agent has not pushed waits on the + // gate ticker; a snapshot that appears mid-wait is picked up on the next + // tick, the chat is pinned to it, and the turn proceeds. + t.Run("WaitsForPush", func(t *testing.T) { + t.Parallel() + fix := newRebindFixture(t) + mClock := quartz.NewMock(t) + trap := mClock.Trap().NewTicker(contextReportTimerTagGroup, contextReportTickerTag) + defer trap.Close() + server := newGateServer(t, fix, mClock, defaultContextReportTimeout) + chat := gateChat(t, fix, fix.agentNoSnap, fix.buildID) + + wc := newGateTurnContext(t, server, chat) + results := awaitGateAsync(fix.ctx, server, wc) + + // The gate found no snapshot and started waiting. + trap.MustWait(fix.ctx).MustRelease(fix.ctx) + + // One empty tick: still nothing to report. + mClock.Advance(contextReportPollInterval).MustWait(fix.ctx) + select { + case res := <-results: + t.Fatalf("gate exited before a snapshot existed: %+v", res) + default: + } + + // The agent's first push arrives mid-wait; the next tick sees it. + hash := []byte{0xcc} + seedAgentContext(fix.ctx, t, fix.db, fix.agentNoSnap, "/home/coder/AGENTS.md", hash, + database.WorkspaceAgentContextBodyKindInstructionFile, + json.RawMessage(`{"instruction_file":{"content":"pushed mid-wait"}}`)) + mClock.Advance(contextReportPollInterval).MustWait(fix.ctx) + + res := testutil.RequireReceive(fix.ctx, t, results) + require.NoError(t, res.err) + require.Equal(t, fix.agentNoSnap, res.agent.ID) + + post, err := fix.db.GetChatByID(fix.ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, hash, post.ContextAggregateHash, "gate exit pins the mid-wait push") + }) + + // CeilingElapses: an agent that never reports fails the turn with a + // visible error naming the wait once the ceiling passes. The ceiling is + // deliberately not a multiple of the poll interval so the timeout fires + // alone. + t.Run("CeilingElapses", func(t *testing.T) { + t.Parallel() + fix := newRebindFixture(t) + mClock := quartz.NewMock(t) + trap := mClock.Trap().NewTicker(contextReportTimerTagGroup, contextReportTickerTag) + defer trap.Close() + ceiling := 2500 * time.Millisecond + server := newGateServer(t, fix, mClock, ceiling) + chat := gateChat(t, fix, fix.agentNoSnap, fix.buildID) + + wc := newGateTurnContext(t, server, chat) + results := awaitGateAsync(fix.ctx, server, wc) + trap.MustWait(fix.ctx).MustRelease(fix.ctx) + + mClock.Advance(contextReportPollInterval).MustWait(fix.ctx) + mClock.Advance(contextReportPollInterval).MustWait(fix.ctx) + mClock.Advance(ceiling - 2*contextReportPollInterval).MustWait(fix.ctx) + + res := testutil.RequireReceive(fix.ctx, t, results) + require.ErrorIs(t, res.err, errTerminalGeneration) + require.ErrorContains(t, res.err, "workspace agent did not report chat context within 2.5s") + }) + + // InterruptDuringWait: canceling the turn context mid-wait propagates + // the cancellation unchanged so the interrupt path sees a plain context + // error, not a terminal gate error that would pollute the chat's error + // state. + t.Run("InterruptDuringWait", func(t *testing.T) { + t.Parallel() + fix := newRebindFixture(t) + mClock := quartz.NewMock(t) + trap := mClock.Trap().NewTicker(contextReportTimerTagGroup, contextReportTickerTag) + defer trap.Close() + server := newGateServer(t, fix, mClock, defaultContextReportTimeout) + chat := gateChat(t, fix, fix.agentNoSnap, fix.buildID) + + turnCtx, cancel := context.WithCancel(fix.ctx) + wc := newGateTurnContext(t, server, chat) + results := awaitGateAsync(turnCtx, server, wc) + trap.MustWait(fix.ctx).MustRelease(fix.ctx) + + cancel() + res := testutil.RequireReceive(fix.ctx, t, results) + require.ErrorIs(t, res.err, context.Canceled) + require.NotErrorIs(t, res.err, errTerminalGeneration) + }) + + // RebindOnNewStartBuild: a pinned chat bound to a previous build's agent + // is rebound to the new start build's agent at turn prep. The rebind + // re-pins, which clears the pinned context because the new agent has not + // pushed, so the gate then waits for the new agent's report and pins the + // chat to it. + t.Run("RebindOnNewStartBuild", func(t *testing.T) { + t.Parallel() + fix := newRebindFixture(t) + mClock := quartz.NewMock(t) + trap := mClock.Trap().NewTicker(contextReportTimerTagGroup, contextReportTickerTag) + defer trap.Close() + server := newGateServer(t, fix, mClock, defaultContextReportTimeout) + + chat := gateChat(t, fix, fix.agentA, fix.buildID) + _, hydrateErr := fix.db.HydrateAgentChatsContext(fix.ctx, database.HydrateAgentChatsContextParams{ + AgentID: fix.agentA, + AggregateHash: fix.hashA, + }) + require.NoError(t, hydrateErr) + newBuild, newAgent := seedStartBuildWithAgent(t, fix.db, fix) + + chat, err := fix.db.GetChatByID(fix.ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, fix.hashA, chat.ContextAggregateHash, "chat starts pinned to the old agent") + + wc := newGateTurnContext(t, server, chat) + results := awaitGateAsync(fix.ctx, server, wc) + trap.MustWait(fix.ctx).MustRelease(fix.ctx) + + // The rebind committed before the gate started waiting: the chat now + // points at the new build's agent and its stale pin is cleared. + mid, err := fix.db.GetChatByID(fix.ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, newAgent.ID, mid.AgentID.UUID, "prep rebinds to the new build's agent") + require.Equal(t, newBuild.ID, mid.BuildID.UUID) + require.Empty(t, mid.ContextAggregateHash, "re-pin clears the old agent's pinned context") + + // Subscribe after the rebind so the received event is the gate-exit + // re-pin's context_ready, not noise from earlier writes. + events := subscribeChatWatchEvents(t, fix) + + hash := []byte{0xdd} + seedAgentContext(fix.ctx, t, fix.db, newAgent.ID, "/home/coder/AGENTS.md", hash, + database.WorkspaceAgentContextBodyKindInstructionFile, + json.RawMessage(`{"instruction_file":{"content":"new build context"}}`)) + mClock.Advance(contextReportPollInterval).MustWait(fix.ctx) + + res := testutil.RequireReceive(fix.ctx, t, results) + require.NoError(t, res.err) + require.Equal(t, newAgent.ID, res.agent.ID) + + post, err := fix.db.GetChatByID(fix.ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, hash, post.ContextAggregateHash, "chat hydrates with the new agent's context") + resources, err := fix.db.ListChatContextResourcesByChatID(fix.ctx, chat.ID) + require.NoError(t, err) + require.Len(t, resources, 1) + require.Equal(t, "/home/coder/AGENTS.md", resources[0].Source) + + // The re-pin branch (empty non-NULL hash) publishes context_ready + // once the pin commits. + event := testutil.RequireReceive(fix.ctx, t, events) + require.Equal(t, codersdk.ChatWatchEventKindContextReady, event.Kind, + "gate-exit re-pinning publishes a context_ready event") + require.Equal(t, chat.ID, event.Chat.ID) + require.NotNil(t, event.Chat.Context) + require.Equal(t, codersdk.ChatContextStateReady, event.Chat.Context.State) + }) +} + +// TestPrepareGenerationContextGate proves the gate end to end through +// prepareGeneration: an unpinned workspace chat blocks until the agent's +// first push, then the prepared turn carries the pushed context (the +// instruction lands in the system prompt inputs and the pushed MCP server +// yields a workspace MCP tool). +func TestPrepareGenerationContextGate(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := chatdTestContext(t) + mClock := quartz.NewMock(t) + trap := mClock.Trap().NewTicker(contextReportTimerTagGroup, contextReportTickerTag) + defer trap.Close() + + user := dbgen.User(t, db, database.User{}) + apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) + org := dbgen.Organization(t, db, database.Organization{}) + dbgen.OrganizationMember(t, db, database.OrganizationMember{ + UserID: user.ID, + OrganizationID: org.ID, + }) + provider := dbgen.AIProviderWithOptionalKey(t, db, database.AIProvider{ + Type: database.AIProviderTypeOpenai, + }, "test-key") + modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Model: "gpt-4o-mini", + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + }, func(p *database.InsertChatModelConfigParams) { + p.Enabled = true + }) + + // A workspace whose agent has not pushed any context yet. + tv := dbgen.TemplateVersion(t, db, database.TemplateVersion{ + OrganizationID: org.ID, + CreatedBy: user.ID, + }) + tmpl := dbgen.Template(t, db, database.Template{ + OrganizationID: org.ID, + ActiveVersionID: tv.ID, + CreatedBy: user.ID, + }) + ws := dbgen.Workspace(t, db, database.WorkspaceTable{ + OwnerID: user.ID, + OrganizationID: org.ID, + TemplateID: tmpl.ID, + }) + pj := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{ + OrganizationID: org.ID, + CompletedAt: sql.NullTime{Valid: true, Time: dbtime.Now()}, + }) + dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{ + WorkspaceID: ws.ID, + TemplateVersionID: tv.ID, + JobID: pj.ID, + Transition: database.WorkspaceTransitionStart, + }) + res := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{ + Transition: database.WorkspaceTransitionStart, + JobID: pj.ID, + }) + agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: res.ID}) + + created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ + OrganizationID: org.ID, + OwnerID: user.ID, + WorkspaceID: uuid.NullUUID{UUID: ws.ID, Valid: true}, + LastModelConfigID: modelConfig.ID, + Title: "context gate end to end", + ClientType: database.ChatClientTypeApi, + InitialMessages: []chatstate.Message{ + { + Role: database.ChatMessageRoleUser, + Content: mustMarshalText(t, "hello"), + Visibility: database.ChatMessageVisibilityBoth, + ModelConfigID: uuid.NullUUID{UUID: modelConfig.ID, Valid: true}, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + ContentVersion: chatprompt.CurrentContentVersion, + APIKeyID: sql.NullString{String: apiKey.ID, Valid: true}, + }, + }, + }) + require.NoError(t, err) + + server := newInternalTestServer( + t, + db, + ps, + chatprovider.ProviderAPIKeys{}, + withInternalTestServerClock(mClock), + withInternalTestServerTransportFactory(&aibridgeTestFactory{}), + ) + + type prepareResult struct { + prepared generationPrepared + err error + } + results := make(chan prepareResult, 1) + go func() { + prepared, err := server.prepareGeneration(ctx, generationPrepareInput{ + Chat: created.Chat, + Messages: created.InitialMessages, + }) + results <- prepareResult{prepared: prepared, err: err} + }() + + // The gate is waiting: the agent has no snapshot yet. + trap.MustWait(ctx).MustRelease(ctx) + + // The agent's first push arrives while the turn waits: an instruction + // file and an MCP server with one tool. + now := dbtime.Now() + hash := []byte{0xab, 0xcd} + _, err = db.UpsertWorkspaceAgentContextSnapshot(ctx, database.UpsertWorkspaceAgentContextSnapshotParams{ + WorkspaceAgentID: agent.ID, + Version: 1, + AggregateHash: hash, + ReceivedAt: now, + }) + require.NoError(t, err) + instructionBody, err := protojson.Marshal(&agentproto.InstructionFileBody{Content: []byte("gate instruction content")}) + require.NoError(t, err) + _, err = db.UpsertWorkspaceAgentContextResource(ctx, database.UpsertWorkspaceAgentContextResourceParams{ + WorkspaceAgentID: agent.ID, + Source: "/home/coder/AGENTS.md", + BodyKind: database.WorkspaceAgentContextBodyKindInstructionFile, + Body: instructionBody, + ContentHash: hash, + SizeBytes: int64(len(instructionBody)), + Status: database.WorkspaceAgentContextResourceStatusOk, + Now: now, + }) + require.NoError(t, err) + schema, err := structpb.NewStruct(map[string]any{ + "type": "object", + "properties": map[string]any{"input": map[string]any{"type": "string"}}, + }) + require.NoError(t, err) + mcpBody, err := protojson.Marshal(&agentproto.MCPServerBody{ + ServerName: "gatesrv", + Tools: []*agentproto.MCPTool{{ + Name: "echo", + Description: "gate echo tool", + InputSchema: schema, + }}, + }) + require.NoError(t, err) + _, err = db.UpsertWorkspaceAgentContextResource(ctx, database.UpsertWorkspaceAgentContextResourceParams{ + WorkspaceAgentID: agent.ID, + Source: "gatesrv", + BodyKind: database.WorkspaceAgentContextBodyKindMcpServer, + Body: mcpBody, + ContentHash: hash, + SizeBytes: int64(len(mcpBody)), + Status: database.WorkspaceAgentContextResourceStatusOk, + Now: now, + }) + require.NoError(t, err) + + mClock.Advance(contextReportPollInterval).MustWait(ctx) + + result := testutil.RequireReceive(ctx, t, results) + require.NoError(t, result.err) + t.Cleanup(result.prepared.Cleanup) + + require.Equal(t, agent.ID, result.prepared.Chat.AgentID.UUID, "turn prep bound the chat to the agent") + + // The pushed MCP server surfaces as a workspace MCP tool on the turn. + require.NotNil(t, findToolByName(result.prepared.Tools, "gatesrv__echo"), + "pushed MCP server must yield a workspace MCP tool") + + // The pushed instruction file lands in the system prompt. + promptText := systemPromptText(t, result.prepared.Prompt) + require.Contains(t, promptText, "gate instruction content", + "pushed instruction file must land in the system prompt") + + post, err := db.GetChatByID(ctx, created.Chat.ID) + require.NoError(t, err) + require.Equal(t, hash, post.ContextAggregateHash, "the turn pinned the pushed snapshot") +} diff --git a/coderd/x/chatd/context_hydration.go b/coderd/x/chatd/context_hydration.go index 147dbf36c75..77de53c2858 100644 --- a/coderd/x/chatd/context_hydration.go +++ b/coderd/x/chatd/context_hydration.go @@ -31,11 +31,11 @@ func latestAgentSnapshot(ctx context.Context, db database.Store, agentID uuid.UU // HydrateAndMarkChatsDirty implements agentapi.ContextDirtyMarker. It runs // inside the PushContextState transaction: it stamps the pushed snapshot hash -// on chats for the agent that have not been hydrated yet (no dirty event), -// then flips already-pinned chats whose hash differs to dirty. It returns a -// callback that publishes the dirty watch events; the caller invokes it only -// after the transaction commits, and the callback is a no-op when nothing -// transitioned to dirty. +// on chats for the agent that have not been hydrated yet, then flips +// already-pinned chats whose hash differs to dirty. It returns a callback +// that publishes the watch events (context_ready for first hydrations, +// context_dirty for drifted chats); the caller invokes it only after the +// transaction commits, and the callback is a no-op when nothing transitioned. // // The pinned hash on dirtied chats is intentionally left unchanged; the // refresh endpoint re-pins it. @@ -44,13 +44,15 @@ func (p *Server) HydrateAndMarkChatsDirty(ctx context.Context, tx database.Store ctx = dbauthz.AsChatd(ctx) // Chats created before the agent's first push land with a NULL pinned - // hash. Stamp them now so they start clean; this is their first - // hydration, so no dirty event is emitted. - if err := tx.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{ + // hash. Stamp them now so they start clean; this first hydration emits + // a ready event (not a dirty one) so watchers flip out of the waiting + // state. + hydrated, err := tx.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{ AgentID: agentID, AggregateHash: aggregateHash, ContextError: snapshotError, - }); err != nil { + }) + if err != nil { return nil, xerrors.Errorf("hydrate agent chats context: %w", err) } @@ -62,15 +64,24 @@ func (p *Server) HydrateAndMarkChatsDirty(ctx context.Context, tx database.Store if err != nil { return nil, xerrors.Errorf("mark chats context dirty: %w", err) } - if len(dirtied) == 0 { + if len(hydrated) == 0 && len(dirtied) == 0 { return func() {}, nil } - // Read the dirtied chats inside the transaction and capture their rows so - // the post-commit callback needs no database access: the published payload - // reflects the just-committed dirty state (no re-read a concurrent refresh - // could race), and the callback does not depend on the request-scoped - // context surviving past commit. Only the transitioned chats are read. + // Read the transitioned chats inside the transaction and capture their + // rows so the post-commit callback needs no database access: the + // published payload reflects the just-committed state (no re-read a + // concurrent refresh could race), and the callback does not depend on + // the request-scoped context surviving past commit. Only the + // transitioned chats are read. + readyChats := make([]database.Chat, 0, len(hydrated)) + for _, id := range hydrated { + chat, err := tx.GetChatByID(ctx, id) + if err != nil { + return nil, xerrors.Errorf("get hydrated chat %s: %w", id, err) + } + readyChats = append(readyChats, chat) + } dirtyChats := make([]database.Chat, 0, len(dirtied)) for _, d := range dirtied { chat, err := tx.GetChatByID(ctx, d.ID) @@ -81,6 +92,7 @@ func (p *Server) HydrateAndMarkChatsDirty(ctx context.Context, tx database.Store } return func() { + p.publishChatPubsubEvents(readyChats, codersdk.ChatWatchEventKindContextReady) p.publishChatPubsubEvents(dirtyChats, codersdk.ChatWatchEventKindContextDirty) }, nil } @@ -104,11 +116,12 @@ func (p *Server) hydrateAgentChatsFromSnapshot(ctx context.Context, agentID uuid if !ok { return nil } - return tx.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{ + _, err = tx.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{ AgentID: agentID, AggregateHash: aggregateHash, ContextError: snapshotError, }) + return err }) } diff --git a/coderd/x/chatd/context_hydration_internal_test.go b/coderd/x/chatd/context_hydration_internal_test.go index 7e71c053e31..dca13482cf0 100644 --- a/coderd/x/chatd/context_hydration_internal_test.go +++ b/coderd/x/chatd/context_hydration_internal_test.go @@ -5,11 +5,14 @@ import ( "testing" "github.com/google/uuid" + "github.com/stretchr/testify/require" "go.uber.org/mock/gomock" "cdr.dev/slog/v3/sloggers/slogtest" "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbmock" + "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" ) @@ -46,7 +49,7 @@ func TestHydrateChatContextOnCreate(t *testing.T) { AgentID: agentID, AggregateHash: snapshot.AggregateHash, ContextError: snapshot.SnapshotError, - }).Return(nil) + }).Return([]uuid.UUID{chat.ID}, nil) server.hydrateChatContextOnCreate(ctx, chat) }) @@ -118,7 +121,7 @@ func TestEnsureChatContextPinnedOnFirstTurn(t *testing.T) { AgentID: agentID, AggregateHash: snapshot.AggregateHash, ContextError: snapshot.SnapshotError, - }).Return(nil) + }).Return([]uuid.UUID{chat.ID}, nil) server.ensureChatContextPinnedOnFirstTurn(ctx, chat) }) @@ -151,3 +154,54 @@ func TestEnsureChatContextPinnedOnFirstTurn(t *testing.T) { server.ensureChatContextPinnedOnFirstTurn(ctx, database.Chat{ID: uuid.New()}) }) } + +// TestHydrateAndMarkChatsDirtyPublishesEvents proves the push-path fan-out +// against a real database: the returned post-commit callback publishes +// context_ready for first hydrations (NULL-hash chats stamped by the push) +// and context_dirty for already-pinned chats whose hash drifted. +func TestHydrateAndMarkChatsDirtyPublishesEvents(t *testing.T) { + t.Parallel() + fix := newRebindFixture(t) + server := &Server{db: fix.db, pubsub: fix.ps, logger: slogtest.Make(t, nil)} + + // A never-hydrated chat (NULL hash) and an already-pinned chat whose + // hash will drift on the push below. + waitingChat := gateChat(t, fix, fix.agentA, fix.buildID) + require.Nil(t, waitingChat.ContextAggregateHash) + pinnedChat := gateChat(t, fix, fix.agentA, fix.buildID) + require.NoError(t, fix.db.SetChatContextSnapshot(fix.ctx, database.SetChatContextSnapshotParams{ + ID: pinnedChat.ID, + AggregateHash: []byte{0x01}, + })) + + events := subscribeChatWatchEvents(t, fix) + + hashNew := []byte{0x77, 0x78} + var publish func() + err := fix.db.InTx(func(tx database.Store) error { + var err error + publish, err = server.HydrateAndMarkChatsDirty( + fix.ctx, tx, fix.agentA, hashNew, "", dbtime.Now()) + return err + }, nil) + require.NoError(t, err) + require.NotNil(t, publish) + publish() + + // One ready event for the hydrated chat and one dirty event for the + // drifted chat; delivery order is not guaranteed, so key by chat. + kindsByChat := map[uuid.UUID]codersdk.ChatWatchEventKind{} + for range 2 { + event := testutil.RequireReceive(fix.ctx, t, events) + kindsByChat[event.Chat.ID] = event.Kind + if event.Chat.ID == waitingChat.ID { + require.NotNil(t, event.Chat.Context) + require.Equal(t, codersdk.ChatContextStateReady, event.Chat.Context.State, + "the ready payload reports the pinned state") + } + } + require.Equal(t, map[uuid.UUID]codersdk.ChatWatchEventKind{ + waitingChat.ID: codersdk.ChatWatchEventKindContextReady, + pinnedChat.ID: codersdk.ChatWatchEventKindContextDirty, + }, kindsByChat) +} diff --git a/coderd/x/chatd/context_integration_test.go b/coderd/x/chatd/context_integration_test.go index 96051bcc5de..75f238c2264 100644 --- a/coderd/x/chatd/context_integration_test.go +++ b/coderd/x/chatd/context_integration_test.go @@ -86,10 +86,14 @@ func TestChatContextDirtyFromAgentPush(t *testing.T) { Status: database.ChatStatusWaiting, }) - // Before any push there is no pinned context. + // Before any push the chat is workspace-bound but unpinned, so the + // context reports the waiting state. got, err := expClient.GetChat(ctx, chat.ID) require.NoError(t, err) - require.Nil(t, got.Context, "no pinned context before the first push") + require.NotNil(t, got.Context, "workspace-bound chats always report context state") + require.Equal(t, codersdk.ChatContextStateWaiting, got.Context.State, + "unpinned workspace chat is waiting for the agent's report") + require.False(t, got.Context.Dirty) requireChatContextNil := func(id uuid.UUID, msg string) { t.Helper() @@ -173,6 +177,8 @@ func TestChatContextDirtyFromAgentPush(t *testing.T) { got, err = expClient.GetChat(ctx, chat.ID) require.NoError(t, err) require.NotNil(t, got.Context, "chat should be hydrated after the initial push") + require.Equal(t, codersdk.ChatContextStateReady, got.Context.State, + "first hydration flips the chat to ready") require.False(t, got.Context.Dirty, "initial hydration is clean") require.Nil(t, got.Context.DirtySince) diff --git a/coderd/x/chatd/context_prompt_internal_test.go b/coderd/x/chatd/context_prompt_internal_test.go index 278d8ea1e17..ff5184ad14e 100644 --- a/coderd/x/chatd/context_prompt_internal_test.go +++ b/coderd/x/chatd/context_prompt_internal_test.go @@ -399,10 +399,11 @@ func TestPinnedWorkspaceContextFromHydratedPin(t *testing.T) { AgentID: uuid.NullUUID{UUID: agent.ID, Valid: true}, Status: database.ChatStatusWaiting, }) - require.NoError(t, db.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{ + _, err := db.HydrateAgentChatsContext(ctx, database.HydrateAgentChatsContextParams{ AgentID: agent.ID, AggregateHash: hash, - })) + }) + require.NoError(t, err) rows, err := db.ListChatContextResourcesByChatID(ctx, chat.ID) require.NoError(t, err) require.Len(t, rows, 2, "the pin holds the agent's instruction file and skill") diff --git a/coderd/x/chatd/context_rebind_internal_test.go b/coderd/x/chatd/context_rebind_internal_test.go index 4c7d62e92ff..e8ccc0425be 100644 --- a/coderd/x/chatd/context_rebind_internal_test.go +++ b/coderd/x/chatd/context_rebind_internal_test.go @@ -18,6 +18,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbmock" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/coderd/database/pubsub" "github.com/coder/coder/v2/testutil" ) @@ -48,10 +49,11 @@ func TestPersistBuildAgentBindingRepinsContext(t *testing.T) { // Pin the chat to agent A through the production hydrate path so it // starts with A's hash and A's resources, exactly as an agent push // would leave it. - require.NoError(t, fix.db.HydrateAgentChatsContext(fix.ctx, database.HydrateAgentChatsContextParams{ + _, hydrateErr := fix.db.HydrateAgentChatsContext(fix.ctx, database.HydrateAgentChatsContextParams{ AgentID: fix.agentA, AggregateHash: fix.hashA, - })) + }) + require.NoError(t, hydrateErr) preRes, err := fix.db.ListChatContextResourcesByChatID(fix.ctx, chat.ID) require.NoError(t, err) require.Len(t, preRes, 1) @@ -121,10 +123,11 @@ func TestPersistBuildAgentBindingRepinsContext(t *testing.T) { AgentID: uuid.NullUUID{UUID: fix.agentA, Valid: true}, Status: database.ChatStatusWaiting, }) - require.NoError(t, fix.db.HydrateAgentChatsContext(fix.ctx, database.HydrateAgentChatsContextParams{ + _, hydrateErr := fix.db.HydrateAgentChatsContext(fix.ctx, database.HydrateAgentChatsContextParams{ AgentID: fix.agentA, AggregateHash: fix.hashA, - })) + }) + require.NoError(t, hydrateErr) preRes, err := fix.db.ListChatContextResourcesByChatID(fix.ctx, chat.ID) require.NoError(t, err) require.Len(t, preRes, 1, "chat starts pinned to agent A") @@ -194,6 +197,7 @@ func TestPersistBuildAgentBindingRepinsContext(t *testing.T) { type rebindFixture struct { db database.Store + ps pubsub.Pubsub ctx context.Context org database.Organization user database.User @@ -217,7 +221,7 @@ type rebindFixture struct { // rebind guard keys on the agent, not the build. func newRebindFixture(t *testing.T) rebindFixture { t.Helper() - db, _ := dbtestutil.NewDB(t) + db, ps := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitLong) user := dbgen.User(t, db, database.User{}) @@ -258,6 +262,7 @@ func newRebindFixture(t *testing.T) rebindFixture { fix := rebindFixture{ db: db, + ps: ps, ctx: ctx, org: org, user: user, diff --git a/coderd/x/chatd/generation_preparer.go b/coderd/x/chatd/generation_preparer.go index 6ca947b1f79..6beb4948312 100644 --- a/coderd/x/chatd/generation_preparer.go +++ b/coderd/x/chatd/generation_preparer.go @@ -223,25 +223,21 @@ func (server *Server) prepareGeneration( promptRows = server.sanitizeForeignProviderExecutedToolRows(ctx, logger, promptRows, modelConfig.ID) if chat.WorkspaceID.Valid { - // Resolve the workspace agent so the chat row's AgentID and - // BuildID bindings are up to date before the chatworker - // decision helper inspects them. ensureWorkspaceAgent does a - // DB lookup and lazily calls persistBuildAgentBinding when - // the bound agent has changed, so this is a cheap metadata - // refresh, not a workspace dial. It must not insert chat - // history; only metadata is mutated here. - agent, _ := workspaceCtx.getWorkspaceAgent(ctx) - - // API-created chats bind their agent lazily here, after - // hydrateChatContextOnCreate ran with no agent. Pin the chat to the - // bound agent's pushed snapshot now if it is still unpinned, so the - // first turn reads workspace context instead of waiting for the - // agent's next push. Idempotent and snapshot-gated; runs before the - // pinned context is read below. - server.ensureChatContextPinnedOnFirstTurn(ctx, workspaceCtx.currentChatSnapshot()) + // A workspace turn must not run before the chat's agent has + // reported its context snapshot. awaitChatContextReported resolves + // the agent (rebinding stale bindings to the latest start build, + // which also re-pins the chat's context), proceeds immediately when + // the chat is already pinned, and otherwise blocks until the bound + // agent has a snapshot, pinning the chat to it before returning. + // Terminal gate failures surface as the chat's visible error state. + agent, gateErr := server.awaitChatContextReported(ctx, &workspaceCtx, logger) + if gateErr != nil { + cleanup() + return generationPrepared{}, gateErr + } var resolveErr error - instruction, workspaceSkills, resolveErr = server.resolveTurnWorkspaceContext(ctx, chat, agent) + instruction, workspaceSkills, resolveErr = server.resolveTurnWorkspaceContext(ctx, workspaceCtx.currentChatSnapshot(), agent) if resolveErr != nil { cleanup() return generationPrepared{}, resolveErr diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index 9ea4d057ad7..8e1b0129172 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -27,6 +27,7 @@ import ( "github.com/coder/coder/v2/coderd/database/dbauthz" "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" + "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/database/pubsub" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/util/ptr" @@ -809,6 +810,17 @@ func seedWorkspaceBinding( JobID: job.ID, }) agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: resource.ID}) + // The context-report gate blocks workspace turns until the agent has + // pushed a context snapshot. Seed an empty snapshot so turns proceed + // as they did before the gate; row existence alone marks the agent as + // having reported, and an empty snapshot contributes no prompt context. + _, err := db.UpsertWorkspaceAgentContextSnapshot(context.Background(), database.UpsertWorkspaceAgentContextSnapshotParams{ + WorkspaceAgentID: agent.ID, + Version: 1, + AggregateHash: []byte{0x01}, + ReceivedAt: dbtime.Now(), + }) + require.NoError(t, err) return workspace, build, agent } diff --git a/codersdk/chats.go b/codersdk/chats.go index 7a02d86aaa5..9d2535f4184 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -141,7 +141,8 @@ type Chat struct { HasUnread bool `json:"has_unread"` // Context reports the chat's pinned workspace-context state and // whether it has drifted from the agent's latest pushed snapshot. - // Nil when the chat has no pinned context yet. + // Present for every workspace-bound chat; nil for chats with no + // workspace and no context remnants. Context *ChatContext `json:"context,omitempty"` Warnings []string `json:"warnings,omitempty"` ClientType ChatClientType `json:"client_type"` @@ -153,10 +154,26 @@ type Chat struct { Children []Chat `json:"children"` } +// ChatContextState describes where a chat stands in the workspace-context +// reporting lifecycle. +type ChatContextState string + +const ( + // ChatContextStateWaiting means the chat is workspace-bound but not + // yet pinned to a reported snapshot; turns gate on the report. + ChatContextStateWaiting ChatContextState = "waiting" + // ChatContextStateReady means the chat is pinned to a reported + // snapshot. + ChatContextStateReady ChatContextState = "ready" +) + // ChatContext reports a chat's pinned workspace context and whether it has // drifted from the agent's latest pushed snapshot. The chat stays usable // when dirty; refreshing re-pins it to the latest snapshot. type ChatContext struct { + // State reports where the chat stands in the context-reporting + // lifecycle: waiting for the agent's first report or ready (pinned). + State ChatContextState `json:"state,omitempty"` // Dirty is true when the agent's latest snapshot hash differs from the // chat's pinned hash. Dirty bool `json:"dirty"` @@ -1817,6 +1834,10 @@ const ( ChatWatchEventKindDeleted ChatWatchEventKind = "deleted" ChatWatchEventKindDiffStatusChange ChatWatchEventKind = "diff_status_change" ChatWatchEventKindActionRequired ChatWatchEventKind = "action_required" + // ChatWatchEventKindContextReady signals that the chat was first + // hydrated or re-pinned from a reported snapshot, so watchers can + // refresh the chat's context state without refetching. + ChatWatchEventKindContextReady ChatWatchEventKind = "context_ready" // ChatWatchEventKindContextDirty signals that the chat's pinned // workspace context drifted from the agent's latest pushed snapshot. // The chat stays usable; a refresh re-pins it to the latest snapshot. diff --git a/docs/reference/api/chats.md b/docs/reference/api/chats.md index 40bbd32637b..06919de04e4 100644 --- a/docs/reference/api/chats.md +++ b/docs/reference/api/chats.md @@ -58,7 +58,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -150,7 +151,7 @@ Status Code **200** | `» build_id` | string(uuid) | false | | | | `» children` | [codersdk.Chat](schemas.md#codersdkchat) | false | | Children holds child (subagent) chats nested under this root chat. Always initialized to an empty slice so the JSON field is present as []. Child chats cannot create their own subagents, so nesting depth is capped at 1 and this slice is always empty for child chats. | | `» client_type` | [codersdk.ChatClientType](schemas.md#codersdkchatclienttype) | false | | | -| `» context` | [codersdk.ChatContext](schemas.md#codersdkchatcontext) | false | | Context reports the chat's pinned workspace-context state and whether it has drifted from the agent's latest pushed snapshot. Nil when the chat has no pinned context yet. | +| `» context` | [codersdk.ChatContext](schemas.md#codersdkchatcontext) | false | | Context reports the chat's pinned workspace-context state and whether it has drifted from the agent's latest pushed snapshot. Present for every workspace-bound chat; nil for chats with no workspace and no context remnants. | | `»» dirty` | boolean | false | | Dirty is true when the agent's latest snapshot hash differs from the chat's pinned hash. | | `»» dirty_since` | string(date-time) | false | | Dirty since is when drift was first detected; nil when not dirty. | | `»» error` | string | false | | Error is the snapshot-level error copied from the pinned snapshot (empty when healthy). | @@ -165,6 +166,7 @@ Status Code **200** | `»»» tools` | array | false | | Tools lists the tools exposed by an MCP server. Populated only for the mcp_server kind; nil otherwise. | | `»»»» description` | string | false | | Description is the tool's human-readable summary; may be empty. | | `»»»» name` | string | false | | Name is the tool name with the "__" prefix the agent adds stripped, so it reads as the server exposes it. | +| `»» state` | [codersdk.ChatContextState](schemas.md#codersdkchatcontextstate) | false | | State reports where the chat stands in the context-reporting lifecycle: waiting for the agent's first report or ready (pinned). | | `» created_at` | string(date-time) | false | | | | `» diff_status` | [codersdk.ChatDiffStatus](schemas.md#codersdkchatdiffstatus) | false | | | | `»» additions` | integer | false | | | @@ -230,6 +232,7 @@ Status Code **200** | `client_type` | `api`, `ui` | | `kind` | `auth`, `config`, `generic`, `instruction_file`, `mcp_config`, `mcp_server`, `missing_key`, `overloaded`, `provider_disabled`, `rate_limit`, `skill`, `stream_silence_timeout`, `timeout`, `usage_limit` | | `status` | `completed`, `error`, `excluded`, `interrupting`, `invalid`, `ok`, `oversize`, `paused`, `pending`, `requires_action`, `running`, `unreadable`, `waiting` | +| `state` | `ready`, `waiting` | | `plan_mode` | `plan` | To perform this operation, you must be authenticated. [Learn more](authentication.md). @@ -333,7 +336,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -426,7 +430,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -676,7 +681,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -823,7 +829,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -916,7 +923,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -1100,7 +1108,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -1193,7 +1202,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -1375,7 +1385,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -1468,7 +1479,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -2219,7 +2231,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -2312,7 +2325,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -2819,7 +2833,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -2912,7 +2927,8 @@ Experimental: this endpoint is subject to change. } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 0d235ae6ba9..6b333421ee5 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -2056,7 +2056,8 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -2149,7 +2150,8 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -2231,7 +2233,7 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in | `build_id` | string | false | | | | `children` | array of [codersdk.Chat](#codersdkchat) | false | | Children holds child (subagent) chats nested under this root chat. Always initialized to an empty slice so the JSON field is present as []. Child chats cannot create their own subagents, so nesting depth is capped at 1 and this slice is always empty for child chats. | | `client_type` | [codersdk.ChatClientType](#codersdkchatclienttype) | false | | | -| `context` | [codersdk.ChatContext](#codersdkchatcontext) | false | | Context reports the chat's pinned workspace-context state and whether it has drifted from the agent's latest pushed snapshot. Nil when the chat has no pinned context yet. | +| `context` | [codersdk.ChatContext](#codersdkchatcontext) | false | | Context reports the chat's pinned workspace-context state and whether it has drifted from the agent's latest pushed snapshot. Present for every workspace-bound chat; nil for chats with no workspace and no context remnants. | | `created_at` | string | false | | | | `diff_status` | [codersdk.ChatDiffStatus](#codersdkchatdiffstatus) | false | | | | `files` | array of [codersdk.ChatFileMetadata](#codersdkchatfilemetadata) | false | | | @@ -2380,7 +2382,8 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in } ] } - ] + ], + "state": "waiting" } ``` @@ -2392,6 +2395,7 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in | `dirty_since` | string | false | | Dirty since is when drift was first detected; nil when not dirty. | | `error` | string | false | | Error is the snapshot-level error copied from the pinned snapshot (empty when healthy). | | `resources` | array of [codersdk.ChatContextResource](#codersdkchatcontextresource) | false | | Resources is the chat's pinned context (instruction files and skills) the prompt is built from, metadata only (no bodies). It is populated only on the single-chat GET response; list and watch payloads leave it nil to stay lightweight. | +| `state` | [codersdk.ChatContextState](#codersdkchatcontextstate) | false | | State reports where the chat stands in the context-reporting lifecycle: waiting for the agent's first report or ready (pinned). | ## codersdk.ChatContextResource @@ -2454,6 +2458,20 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in |-------------------------------------------------------| | `excluded`, `invalid`, `ok`, `oversize`, `unreadable` | +## codersdk.ChatContextState + +```json +"waiting" +``` + +### Properties + +#### Enumerated Values + +| Value(s) | +|--------------------| +| `ready`, `waiting` | + ## codersdk.ChatContextTool ```json @@ -3934,7 +3952,8 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in } ] } - ] + ], + "state": "waiting" }, "created_at": "2019-08-24T14:15:22Z", "diff_status": { @@ -4034,9 +4053,9 @@ AuthorizationObject can represent a "set" of objects, such as: all workspaces in #### Enumerated Values -| Value(s) | -|-----------------------------------------------------------------------------------------------------------------------------------| -| `action_required`, `context_dirty`, `created`, `deleted`, `diff_status_change`, `status_change`, `summary_change`, `title_change` | +| Value(s) | +|----------------------------------------------------------------------------------------------------------------------------------------------------| +| `action_required`, `context_dirty`, `context_ready`, `created`, `deleted`, `diff_status_change`, `status_change`, `summary_change`, `title_change` | ## codersdk.ClusterConfig diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 7b1a948d8d8..798e6113338 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -2198,6 +2198,49 @@ describe("mergeWatchedChatSummary", () => { }); }); + it("applies context_ready state while preserving the pinned resource list", () => { + const cachedChat = makeChat("chat-1", { + updated_at: "2025-01-01T00:00:00.000Z", + context: { + state: "waiting", + dirty: false, + resources: [ + { + source: "/AGENTS.md", + kind: "instruction_file", + size_bytes: 10, + status: "ok", + }, + ], + }, + }); + const watchedChat = makeChat("chat-1", { + // Readiness is tracked outside updated_at, so an older event + // timestamp still applies the state flip. + updated_at: "2024-12-31T00:00:00.000Z", + context: { state: "ready", dirty: false }, + }); + + expect( + mergeWatchedChatSummary(cachedChat, watchedChat, { + eventKind: "context_ready", + }).context, + ).toEqual({ + state: "ready", + dirty: false, + // The lightweight watch payload omits resources; the merge keeps the + // pinned list a prior single-chat GET populated. + resources: [ + { + source: "/AGENTS.md", + kind: "instruction_file", + size_bytes: 10, + status: "ok", + }, + ], + }); + }); + it("leaves context untouched for non-context events", () => { const context = { dirty: true, dirty_since: "2025-01-02T00:00:00.000Z" }; const cachedChat = makeChat("chat-1", { diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 8262c8f4918..6157bfa6eb8 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -422,6 +422,7 @@ export const mergeWatchedChatSummary = ( const isSummaryEvent = eventKind === "summary_change"; const isDiffStatusEvent = eventKind === "diff_status_change"; const isContextDirtyEvent = eventKind === "context_dirty"; + const isContextReadyEvent = eventKind === "context_ready"; const updatedAtComparison = compareUpdatedAtInstants( cachedChat.updated_at, watchedChat.updated_at, @@ -437,13 +438,14 @@ export const mergeWatchedChatSummary = ( const nextDiffStatus = isDiffStatusEvent ? watchedChat.diff_status : cachedChat.diff_status; - // Context drift is tracked outside chats.updated_at (it is driven by - // agent context pushes), so apply context_dirty payloads regardless of - // the summary timestamp. Merge rather than replace so the pinned - // resources a single-chat GET populated are preserved while the dirty - // flags update; the open chat refetches the full detail. + // Context drift and readiness are tracked outside chats.updated_at + // (they are driven by agent context pushes), so apply context_dirty + // and context_ready payloads regardless of the summary timestamp. + // Merge rather than replace so the pinned resources a single-chat GET + // populated are preserved while the state flags update; the open chat + // refetches the full detail. const nextContext = - isContextDirtyEvent && watchedChat.context + (isContextDirtyEvent || isContextReadyEvent) && watchedChat.context ? { ...cachedChat.context, ...watchedChat.context } : cachedChat.context; const nextWorkspaceId = isFreshEnough diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 740a5e729fe..521702d17db 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -1634,7 +1634,8 @@ export interface Chat { /** * Context reports the chat's pinned workspace-context state and * whether it has drifted from the agent's latest pushed snapshot. - * Nil when the chat has no pinned context yet. + * Present for every workspace-bound chat; nil for chats with no + * workspace and no context remnants. */ readonly context?: ChatContext; readonly warnings?: readonly string[]; @@ -1732,6 +1733,11 @@ export interface ChatConfig { * when dirty; refreshing re-pins it to the latest snapshot. */ export interface ChatContext { + /** + * State reports where the chat stands in the context-reporting + * lifecycle: waiting for the agent's first report or ready (pinned). + */ + readonly state?: ChatContextState; /** * Dirty is true when the agent's latest snapshot hash differs from the * chat's pinned hash. @@ -1849,6 +1855,11 @@ export const ChatContextResourceStatuses: ChatContextResourceStatus[] = [ "unreadable", ]; +// From codersdk/chats.go +export type ChatContextState = "ready" | "waiting"; + +export const ChatContextStates: ChatContextState[] = ["ready", "waiting"]; + // From codersdk/chats.go /** * ChatContextTool is one tool exposed by a pinned MCP server, reported on the @@ -3362,6 +3373,7 @@ export interface ChatWatchEvent { export type ChatWatchEventKind = | "action_required" | "context_dirty" + | "context_ready" | "created" | "deleted" | "diff_status_change" @@ -3372,6 +3384,7 @@ export type ChatWatchEventKind = export const ChatWatchEventKinds: ChatWatchEventKind[] = [ "action_required", "context_dirty", + "context_ready", "created", "deleted", "diff_status_change", diff --git a/site/src/pages/AgentsPage/AgentsPage.tsx b/site/src/pages/AgentsPage/AgentsPage.tsx index b706153711f..685404144d4 100644 --- a/site/src/pages/AgentsPage/AgentsPage.tsx +++ b/site/src/pages/AgentsPage/AgentsPage.tsx @@ -617,7 +617,10 @@ const AgentsPage: FC = () => { if (shouldInvalidateFilteredChatList(updatedChat, chatEvent.kind)) { void invalidateChatListQueries(queryClient); } - if (chatEvent.kind === "context_dirty") { + if ( + chatEvent.kind === "context_dirty" || + chatEvent.kind === "context_ready" + ) { // The watch payload carries only the lightweight // context flags (the merge above applies them); // refetch the open chat to pull the pinned diff --git a/site/src/pages/AgentsPage/components/ContextUsageIndicator.stories.tsx b/site/src/pages/AgentsPage/components/ContextUsageIndicator.stories.tsx index 22b8c3ed58a..6245f4a499b 100644 --- a/site/src/pages/AgentsPage/components/ContextUsageIndicator.stories.tsx +++ b/site/src/pages/AgentsPage/components/ContextUsageIndicator.stories.tsx @@ -3,6 +3,7 @@ import { expect, fn, userEvent, waitFor, within } from "storybook/test"; import { MockChatContextClean, MockChatContextDirty, + MockChatContextWaiting, } from "#/testHelpers/chatEntities"; import { ContextUsageIndicator } from "./ContextUsageIndicator"; @@ -204,6 +205,31 @@ export const Dirty: Story = { }, }; +// Waiting pin: the chat is workspace-bound but the agent has not reported a +// context snapshot yet, so the indicator shows a subtle spinner and the +// popover explains the wait. No refresh affordance: there is nothing to +// re-pin yet. +export const Waiting: Story = { + args: { + usage: { + context: MockChatContextWaiting, + }, + }, + play: async ({ canvasElement }) => { + const button = within(canvasElement).getByRole("button"); + expect(button.getAttribute("aria-label") ?? "").toContain( + "Waiting for workspace context", + ); + + await userEvent.hover(button); + const body = within(document.body); + await waitFor(() => + expect(body.getByText("Waiting for workspace context")).toBeVisible(), + ); + expect(body.queryByRole("button", { name: "Refresh context" })).toBeNull(); + }, +}; + // Snapshot-level error: the ring shows a distinct error treatment and the // popover surfaces the error message. export const SnapshotError: Story = { diff --git a/site/src/pages/AgentsPage/components/ContextUsageIndicator.tsx b/site/src/pages/AgentsPage/components/ContextUsageIndicator.tsx index 0111490ef4d..a570c1ef83f 100644 --- a/site/src/pages/AgentsPage/components/ContextUsageIndicator.tsx +++ b/site/src/pages/AgentsPage/components/ContextUsageIndicator.tsx @@ -238,6 +238,9 @@ export const ContextUsageIndicator: FC<{ const context = usage?.context; const isDirty = context?.dirty ?? false; + // A waiting chat is workspace-bound but not yet pinned to a reported + // snapshot; turns gate on the agent's first context report. + const isWaiting = context?.state === "waiting"; const contextError = context?.error ?? ""; const hasContextError = contextError !== ""; const pinnedResources = context?.resources; @@ -327,14 +330,23 @@ export const ContextUsageIndicator: FC<{ const fileGroups = groupByDirectory(fileItems); const skillGroups = groupByDirectory(skillItems); - const ariaLabel = hasPercent - ? `Context usage ${percentLabel}. ${formatTokenCount(usedTokens)} of ${formatTokenCount(contextLimitTokens)} tokens used.${isDirty ? " Context changed." : ""}` + const ariaStateSuffix = isWaiting + ? " Waiting for workspace context." : isDirty - ? "Context usage. Context changed." - : "Context usage"; + ? " Context changed." + : ""; + const ariaLabel = hasPercent + ? `Context usage ${percentLabel}. ${formatTokenCount(usedTokens)} of ${formatTokenCount(contextLimitTokens)} tokens used.${ariaStateSuffix}` + : `Context usage${ariaStateSuffix ? `.${ariaStateSuffix}` : ""}`; const panelContent = (
+ {isWaiting && ( +
+ + Waiting for workspace context +
+ )} {hasPercent ? `${percentLabel} - ${formatTokenCountCompact(usedTokens)} / ${formatTokenCountCompact(contextLimitTokens)} context used` : "Context usage unavailable"} @@ -570,6 +582,14 @@ export const ContextUsageIndicator: FC<{ progressClassName="stroke-current" className={cn("size-icon-sm", toneClassName)} /> + {isWaiting && ( + + )} {(isDirty || hasContextError) && ( Date: Sun, 12 Jul 2026 21:50:44 +0000 Subject: [PATCH 2/2] test(coderd/database/dbauthz): update HydrateAgentChatsContext mock for hydrated IDs --- coderd/database/dbauthz/dbauthz_test.go | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 67c7cee24c6..5caf6090cc3 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -558,8 +558,9 @@ func (s *MethodTestSuite) TestChats() { })) s.Run("HydrateAgentChatsContext", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { arg := database.HydrateAgentChatsContextParams{AgentID: uuid.New()} - dbm.EXPECT().HydrateAgentChatsContext(gomock.Any(), arg).Return(nil).AnyTimes() - check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate) + hydrated := []uuid.UUID{uuid.New()} + dbm.EXPECT().HydrateAgentChatsContext(gomock.Any(), arg).Return(hydrated, nil).AnyTimes() + check.Args(arg).Asserts(rbac.ResourceChat, policy.ActionUpdate).Returns(hydrated) })) s.Run("MarkChatsContextDirtyByAgent", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { arg := database.MarkChatsContextDirtyByAgentParams{AgentID: uuid.New()}