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/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()} 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) && (