diff --git a/cli/exp_chat.go b/cli/exp_chat.go index 61c017f172..396504d658 100644 --- a/cli/exp_chat.go +++ b/cli/exp_chat.go @@ -1,14 +1,20 @@ package cli import ( + "context" "fmt" "os" + "path" "path/filepath" + "strings" "github.com/google/uuid" "golang.org/x/xerrors" "github.com/coder/coder/v2/agent/agentcontextconfig" + "github.com/coder/coder/v2/agent/agentsocket" + "github.com/coder/coder/v2/cli/cliui" + "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/agentsdk" "github.com/coder/serpent" ) @@ -28,102 +34,372 @@ func (r *RootCmd) chatCommand() *serpent.Command { } func (r *RootCmd) chatContextCommand() *serpent.Command { + // socketPath is shared by the in-workspace source commands (list, show, + // add, remove) and the no-argument refresh, which all talk to the agent's + // local IPC socket. + var socketPath string return &serpent.Command{ Use: "context", - Short: "Manage chat context", - Long: "Add or clear context files and skills for an active chat session.", + Short: "Manage workspace context", + Long: "Inspect and manage the workspace context sources (instruction files, " + + "skills, and MCP configs) the agent resolves, and refresh a chat to the " + + "agent's latest snapshot.\n\nThe list, show, add, and remove commands manage " + + "agent-local sources and must be run from inside the workspace.", Handler: func(i *serpent.Invocation) error { return i.Command.HelpHandler(i) }, Children: []*serpent.Command{ - r.chatContextAddCommand(), + r.chatContextListCommand(&socketPath), + r.chatContextShowCommand(&socketPath), + r.chatContextAddCommand(&socketPath), + r.chatContextRemoveCommand(&socketPath), + r.chatContextRefreshCommand(&socketPath), r.chatContextClearCommand(), }, + Options: serpent.OptionSet{{ + Flag: "socket-path", + Env: "CODER_AGENT_SOCKET_PATH", + Description: "Path to the agent socket used by the in-workspace source commands.", + Value: serpent.StringOf(&socketPath), + }}, } } -func (*RootCmd) chatContextAddCommand() *serpent.Command { - var ( - dir string - chatID string +// resolveContextSourcePath makes a user-supplied source path absolute so the +// agent (which requires absolute, canonical paths) accepts it. A leading ~ is +// preserved for the agent to expand against its own home directory. A path that +// is already absolute on the agent's POSIX filesystem (a leading /) is cleaned +// and passed through; filepath.Abs is host-OS specific and would mangle such a +// path on a Windows CLI host, so it is reserved for resolving relative paths +// against the CLI's working directory, which shares the workspace filesystem +// with the agent. +func resolveContextSourcePath(p string) (string, error) { + p = strings.TrimSpace(p) + if p == "" { + return "", xerrors.New("path is empty") + } + if p == "~" || strings.HasPrefix(p, "~/") { + return p, nil + } + if strings.HasPrefix(p, "/") { + return path.Clean(p), nil + } + abs, err := filepath.Abs(p) + if err != nil { + return "", xerrors.Errorf("resolve path %q: %w", p, err) + } + return abs, nil +} + +// dialAgentContextSocket connects to the workspace agent's local IPC socket. +// It is only reachable from inside the workspace. +func dialAgentContextSocket(ctx context.Context, socketPath string) (*agentsocket.Client, error) { + opts := []agentsocket.Option{} + if socketPath != "" { + opts = append(opts, agentsocket.WithPath(socketPath)) + } + client, err := agentsocket.NewClient(ctx, opts...) + if err != nil { + return nil, xerrors.Errorf("connect to agent socket (run this from inside the workspace): %w", err) + } + return client, nil +} + +func (*RootCmd) chatContextListCommand(socketPath *string) *serpent.Command { + formatter := cliui.NewOutputFormatter( + cliui.TableFormat([]agentsocket.ContextSource{}, []string{"path"}), + cliui.JSONFormat(), ) + cmd := &serpent.Command{ + Use: "list", + Short: "List the workspace context sources registered on the agent", + Long: "List the additional scan roots registered on this workspace's agent. " + + "Built-in defaults (the working directory, ~/.coder, ~/.claude) are always " + + "scanned and are not shown here.\n\nMust be run from inside the workspace.", + Middleware: serpent.RequireNArgs(0), + Handler: func(inv *serpent.Invocation) error { + ctx := inv.Context() + client, err := dialAgentContextSocket(ctx, *socketPath) + if err != nil { + return err + } + defer client.Close() + + sources, err := client.ContextSources(ctx) + if err != nil { + return xerrors.Errorf("list context sources: %w", err) + } + if len(sources) == 0 && formatter.FormatID() == "table" { + cliui.Info(inv.Stdout, "No context sources registered.") + return nil + } + out, err := formatter.Format(ctx, sources) + if err != nil { + return xerrors.Errorf("format output: %w", err) + } + _, _ = fmt.Fprintln(inv.Stdout, out) + return nil + }, + } + formatter.AttachOptions(&cmd.Options) + return cmd +} + +func (*RootCmd) chatContextShowCommand(socketPath *string) *serpent.Command { + formatter := cliui.NewOutputFormatter( + cliui.TableFormat( + []agentsocket.ContextResource{}, + []string{"kind", "name", "source", "status", "size bytes", "error"}, + ), + cliui.JSONFormat(), + ) + cmd := &serpent.Command{ + Use: "show ", + Short: "Show a context source and the resources it contributes", + Long: "Show a registered context source and the resources the agent currently " + + "resolves from it (instruction files, skills, MCP configs), including any " + + "that failed to read or parse.\n\nMust be run from inside the workspace.", + Middleware: serpent.RequireNArgs(1), + Handler: func(inv *serpent.Invocation) error { + ctx := inv.Context() + client, err := dialAgentContextSocket(ctx, *socketPath) + if err != nil { + return err + } + defer client.Close() + + path, err := resolveContextSourcePath(inv.Args[0]) + if err != nil { + return err + } + src, err := client.GetContextSource(ctx, path) + if err != nil { + return xerrors.Errorf("get context source: %w", err) + } + snap, err := client.GetContextSnapshot(ctx) + if err != nil { + return xerrors.Errorf("get context snapshot: %w", err) + } + resources := make([]agentsocket.ContextResource, 0, len(snap.Resources)) + for _, res := range snap.Resources { + if res.SourcePath == src.Path { + resources = append(resources, res) + } + } + + if formatter.FormatID() == "table" { + cliui.Infof(inv.Stdout, "Source: %s (%d resources)", src.Path, len(resources)) + } + out, err := formatter.Format(ctx, resources) + if err != nil { + return xerrors.Errorf("format output: %w", err) + } + _, _ = fmt.Fprintln(inv.Stdout, out) + return nil + }, + } + formatter.AttachOptions(&cmd.Options) + return cmd +} + +func (*RootCmd) chatContextAddCommand(socketPath *string) *serpent.Command { + var chatID string agentAuth := &AgentAuth{} cmd := &serpent.Command{ - Use: "add", - Short: "Add context to an active chat", - Long: "Read instruction files and discover skills from a directory, then add " + - "them as context to an active chat session. Multiple calls " + - "are additive.", + Use: "add ", + Short: "Register a workspace context source", + Long: "Register a path as an additional context source on this workspace's agent. " + + "The agent treats it as an extra scan root, applying the same discovery rules " + + "it uses for the working directory: AGENTS.md / CLAUDE.md / .cursorrules, " + + ".agents/skills//SKILL.md, and .mcp.json are picked up now and as they " + + "appear. Any change to a recognized file dirties this workspace's chats until " + + "you refresh.\n\nA path may be a file or a directory. Must be run from inside " + + "the workspace.\n\nPass --chat to keep the legacy one-shot behavior: read " + + "context from the path once and inject it into a single chat without " + + "registering a source.", + Middleware: serpent.RequireNArgs(1), Handler: func(inv *serpent.Invocation) error { ctx := inv.Context() ctx, stop := inv.SignalNotifyContext(ctx, StopSignals...) defer stop() - if dir == "" && inv.Environ.Get("CODER") != "true" { - return xerrors.New("this command must be run inside a Coder workspace (set --dir to override)") + // Legacy one-shot inject into a single chat. + if chatID != "" { + return addChatContextOneShot(ctx, inv, agentAuth, inv.Args[0], chatID) } - client, err := agentAuth.CreateClient() - if err != nil { - return xerrors.Errorf("create agent client: %w", err) - } - - resolvedDir := dir - if resolvedDir == "" { - resolvedDir, err = os.Getwd() - if err != nil { - return xerrors.Errorf("get working directory: %w", err) - } - } - resolvedDir, err = filepath.Abs(resolvedDir) - if err != nil { - return xerrors.Errorf("resolve directory: %w", err) - } - info, err := os.Stat(resolvedDir) - if err != nil { - return xerrors.Errorf("cannot read directory %q: %w", resolvedDir, err) - } - if !info.IsDir() { - return xerrors.Errorf("%q is not a directory", resolvedDir) - } - - parts := agentcontextconfig.ContextPartsFromDir(resolvedDir) - if len(parts) == 0 { - _, _ = fmt.Fprintln(inv.Stderr, "No context files or skills found in "+resolvedDir) - return nil - } - - // Resolve chat ID from flag or auto-detect. - resolvedChatID, err := parseChatID(chatID) + // Source registration (default). + path, err := resolveContextSourcePath(inv.Args[0]) if err != nil { return err } - - resp, err := client.AddChatContext(ctx, agentsdk.AddChatContextRequest{ - ChatID: resolvedChatID, - Parts: parts, - }) + client, err := dialAgentContextSocket(ctx, *socketPath) if err != nil { - return xerrors.Errorf("add chat context: %w", err) + return err } + defer client.Close() - _, _ = fmt.Fprintf(inv.Stdout, "Added %d context part(s) to chat %s\n", resp.Count, resp.ChatID) + src, err := client.AddContextSource(ctx, path) + if err != nil { + return xerrors.Errorf("add context source: %w", err) + } + _, _ = fmt.Fprintf(inv.Stdout, "Registered context source %s\n", src.Path) return nil }, - Options: serpent.OptionSet{ - { - Name: "Directory", - Flag: "dir", - Description: "Directory to read context files and skills from. Defaults to the current working directory.", - Value: serpent.StringOf(&dir), - }, - { - Name: "Chat ID", - Flag: "chat", - Env: "CODER_CHAT_ID", - Description: "Chat ID to add context to. Auto-detected from CODER_CHAT_ID, the only active chat, or the only top-level active chat.", - Value: serpent.StringOf(&chatID), - }, + Options: serpent.OptionSet{{ + Name: "Chat ID", + Flag: "chat", + Env: "CODER_CHAT_ID", + Description: "Inject context from into a single chat (legacy one-shot) instead of registering a source. Auto-detected from CODER_CHAT_ID, the only active chat, or the only top-level active chat.", + Value: serpent.StringOf(&chatID), + }}, + } + agentAuth.AttachOptions(cmd, false) + return cmd +} + +// addChatContextOneShot preserves the legacy `add --chat` behavior: read +// context files and skills from a directory and inject them into a single +// chat via coderd, without registering a persistent source. +func addChatContextOneShot(ctx context.Context, inv *serpent.Invocation, agentAuth *AgentAuth, dir, chatID string) error { + client, err := agentAuth.CreateClient() + if err != nil { + return xerrors.Errorf("create agent client: %w", err) + } + + resolvedDir, err := filepath.Abs(dir) + if err != nil { + return xerrors.Errorf("resolve directory: %w", err) + } + info, err := os.Stat(resolvedDir) + if err != nil { + return xerrors.Errorf("cannot read directory %q: %w", resolvedDir, err) + } + if !info.IsDir() { + return xerrors.Errorf("--chat one-shot inject requires a directory, but %q is a file", resolvedDir) + } + + parts := agentcontextconfig.ContextPartsFromDir(resolvedDir) + if len(parts) == 0 { + _, _ = fmt.Fprintln(inv.Stderr, "No context files or skills found in "+resolvedDir) + return nil + } + + resolvedChatID, err := parseChatID(chatID) + if err != nil { + return err + } + + resp, err := client.AddChatContext(ctx, agentsdk.AddChatContextRequest{ + ChatID: resolvedChatID, + Parts: parts, + }) + if err != nil { + return xerrors.Errorf("add chat context: %w", err) + } + + _, _ = fmt.Fprintf(inv.Stdout, "Added %d context part(s) to chat %s\n", resp.Count, resp.ChatID) + return nil +} + +func (*RootCmd) chatContextRemoveCommand(socketPath *string) *serpent.Command { + cmd := &serpent.Command{ + Use: "remove ", + Short: "Remove a workspace context source", + Long: "Remove a previously-registered context source from this workspace's agent " + + "and re-resolve. Built-in default scan roots cannot be removed.\n\nMust be run " + + "from inside the workspace.", + Middleware: serpent.RequireNArgs(1), + Handler: func(inv *serpent.Invocation) error { + ctx := inv.Context() + client, err := dialAgentContextSocket(ctx, *socketPath) + if err != nil { + return err + } + defer client.Close() + + path, err := resolveContextSourcePath(inv.Args[0]) + if err != nil { + return err + } + if err := client.RemoveContextSource(ctx, path); err != nil { + return xerrors.Errorf("remove context source: %w", err) + } + _, _ = fmt.Fprintf(inv.Stdout, "Removed context source %s\n", path) + return nil + }, + } + return cmd +} + +func (r *RootCmd) chatContextRefreshCommand(socketPath *string) *serpent.Command { + agentAuth := &AgentAuth{} + cmd := &serpent.Command{ + Use: "refresh []", + Short: "Refresh chat context to the agent's latest snapshot", + Long: "Re-pin a chat to the workspace agent's latest context snapshot and clear " + + "its drift marker. The chat's next turn uses the refreshed context.\n\nWith a " + + " argument, refreshes that chat and works from anywhere.\n\nWith no " + + "argument, run from inside the workspace: forces the agent to re-resolve its " + + "sources (catching freshly-cloned repos and startup-script writes the watcher " + + "has not seen yet), then refreshes every drifted chat. This path authenticates " + + "with the agent token, so it does not require 'coder login'.", + Middleware: serpent.RequireRangeArgs(0, 1), + Handler: func(inv *serpent.Invocation) error { + ctx := inv.Context() + + // With a argument: refresh that specific chat through the + // user-facing API. Works from anywhere with a logged-in CLI. + if len(inv.Args) == 1 { + chatID, err := uuid.Parse(inv.Args[0]) + if err != nil { + return xerrors.Errorf("invalid chat ID %q: %w", inv.Args[0], err) + } + client, err := r.InitClient(inv) + if err != nil { + return err + } + exp := codersdk.NewExperimentalClient(client) + chat, err := exp.RefreshChatContext(ctx, chatID) + if err != nil { + return xerrors.Errorf("refresh chat context: %w", err) + } + _, _ = fmt.Fprintf(inv.Stdout, "Refreshed context for chat %s.\n", chatID) + if chat.Context != nil && chat.Context.Error != "" { + _, _ = fmt.Fprintf(inv.Stdout, "Snapshot reported an error: %s\n", chat.Context.Error) + } + return nil + } + + // No argument: in-workspace. Re-resolve the agent's sources over + // the local context socket, then ask the agent (using its own + // token) to re-pin every drifted chat. Neither step needs a + // logged-in user session. + sock, err := dialAgentContextSocket(ctx, *socketPath) + if err != nil { + return xerrors.Errorf("connect to agent context socket "+ + "(run inside the workspace, or pass a ID): %w", err) + } + defer sock.Close() + snap, err := sock.ResyncContext(ctx) + if err != nil { + return xerrors.Errorf("re-resolve agent context: %w", err) + } + _, _ = fmt.Fprintf(inv.Stdout, "Re-resolved agent context (version %d, %d resources).\n", snap.Version, len(snap.Resources)) + if snap.SnapshotError != "" { + _, _ = fmt.Fprintf(inv.Stdout, "Snapshot reported an error: %s\n", snap.SnapshotError) + } + + agentClient, err := agentAuth.CreateClient() + if err != nil { + return xerrors.Errorf("create agent client: %w", err) + } + resp, err := agentClient.RefreshChatContext(ctx) + if err != nil { + return xerrors.Errorf("refresh chat context: %w", err) + } + _, _ = fmt.Fprintf(inv.Stdout, "Refreshed %d drifted chat(s).\n", resp.Refreshed) + return nil }, } agentAuth.AttachOptions(cmd, false) diff --git a/cli/exp_chat_internal_test.go b/cli/exp_chat_internal_test.go new file mode 100644 index 0000000000..68dbec0e7e --- /dev/null +++ b/cli/exp_chat_internal_test.go @@ -0,0 +1,77 @@ +package cli + +import ( + "path/filepath" + "testing" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" +) + +func TestParseChatID(t *testing.T) { + t.Parallel() + + t.Run("EmptyIsNil", func(t *testing.T) { + t.Parallel() + got, err := parseChatID("") + require.NoError(t, err) + require.Equal(t, uuid.Nil, got) + }) + + t.Run("ValidUUID", func(t *testing.T) { + t.Parallel() + want := uuid.MustParse("11111111-1111-4111-8111-111111111111") + got, err := parseChatID(want.String()) + require.NoError(t, err) + require.Equal(t, want, got) + }) + + t.Run("InvalidErrors", func(t *testing.T) { + t.Parallel() + _, err := parseChatID("not-a-uuid") + require.Error(t, err) + require.Contains(t, err.Error(), "invalid chat ID") + }) +} + +func TestResolveContextSourcePath(t *testing.T) { + t.Parallel() + + t.Run("EmptyErrors", func(t *testing.T) { + t.Parallel() + _, err := resolveContextSourcePath(" ") + require.Error(t, err) + require.Contains(t, err.Error(), "empty") + }) + + t.Run("PreservesTilde", func(t *testing.T) { + t.Parallel() + // A leading ~ is left for the agent to expand against its own home. + got, err := resolveContextSourcePath("~") + require.NoError(t, err) + require.Equal(t, "~", got) + + got, err = resolveContextSourcePath(" ~/skills/deploy ") + require.NoError(t, err) + require.Equal(t, "~/skills/deploy", got) + }) + + t.Run("KeepsAbsolute", func(t *testing.T) { + t.Parallel() + got, err := resolveContextSourcePath("/home/coder/AGENTS.md") + require.NoError(t, err) + require.Equal(t, "/home/coder/AGENTS.md", got) + }) + + t.Run("MakesRelativeAbsolute", func(t *testing.T) { + t.Parallel() + // "./" was the reported failure: a relative path must be resolved to an + // absolute one before it reaches the agent. + got, err := resolveContextSourcePath("./") + require.NoError(t, err) + require.True(t, filepath.IsAbs(got), "want absolute, got %q", got) + want, err := filepath.Abs("./") + require.NoError(t, err) + require.Equal(t, want, got) + }) +} diff --git a/cli/exp_chat_test.go b/cli/exp_chat_test.go index 30696c6eca..f204db3301 100644 --- a/cli/exp_chat_test.go +++ b/cli/exp_chat_test.go @@ -1,6 +1,7 @@ package cli_test import ( + "path/filepath" "testing" "github.com/stretchr/testify/require" @@ -11,36 +12,31 @@ import ( func TestExpChatContextAdd(t *testing.T) { t.Parallel() - t.Run("RequiresWorkspaceOrDir", func(t *testing.T) { + t.Run("RequiresPathArgument", func(t *testing.T) { t.Parallel() + // `add` registers a context source identified by , so the path + // argument is required and a bare invocation is a usage error. inv, _ := clitest.New(t, "exp", "chat", "context", "add") err := inv.Run() require.Error(t, err) - require.Contains(t, err.Error(), "this command must be run inside a Coder workspace") + require.Contains(t, err.Error(), "wanted 1 args but got 0") }) - t.Run("AllowsExplicitDir", func(t *testing.T) { + t.Run("RequiresWorkspaceSocket", func(t *testing.T) { t.Parallel() - inv, _ := clitest.New(t, "exp", "chat", "context", "add", "--dir", t.TempDir()) + // Source registration talks to the agent over its local socket, so + // outside a workspace it fails to connect rather than silently doing + // nothing. Point at a socket path that does not exist so the dial + // fails deterministically (and never touches a real agent socket). + missingSocket := filepath.Join(t.TempDir(), "agent.sock") + inv, _ := clitest.New(t, "exp", "chat", "context", "add", t.TempDir(), + "--socket-path", missingSocket) err := inv.Run() - if err != nil { - require.NotContains(t, err.Error(), "this command must be run inside a Coder workspace") - } - }) - - t.Run("AllowsWorkspaceEnv", func(t *testing.T) { - t.Parallel() - - inv, _ := clitest.New(t, "exp", "chat", "context", "add") - inv.Environ.Set("CODER", "true") - - err := inv.Run() - if err != nil { - require.NotContains(t, err.Error(), "this command must be run inside a Coder workspace") - } + require.Error(t, err) + require.Contains(t, err.Error(), "inside the workspace") }) } diff --git a/coderd/coderd.go b/coderd/coderd.go index c9a76cbab7..1326895ce5 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -1796,6 +1796,7 @@ func New(options *Options) *API { r.Route("/experimental", func(r chi.Router) { r.Post("/chat-context", api.workspaceAgentAddChatContext) r.Delete("/chat-context", api.workspaceAgentClearChatContext) + r.Post("/chat-context/refresh", api.workspaceAgentRefreshChatContext) }) r.Route("/tasks/{task}", func(r chi.Router) { r.Post("/log-snapshot", api.postWorkspaceAgentTaskLogSnapshot) diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go index 0dc91010cc..f4c0aa0c99 100644 --- a/coderd/workspaceagents.go +++ b/coderd/workspaceagents.go @@ -2691,6 +2691,69 @@ func (api *API) workspaceAgentClearChatContext(rw http.ResponseWriter, r *http.R }) } +// workspaceAgentRefreshChatContext re-pins every drifted chat bound to the +// calling agent to the agent's latest context snapshot, clearing their +// drift markers. It backs the in-workspace `coder exp chat context refresh` +// (no chat argument), which uses the agent token rather than a user +// session, mirroring workspaceAgentClearChatContext's auth model. +// +// @x-apidocgen {"skip": true} +func (api *API) workspaceAgentRefreshChatContext(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + workspaceAgent := httpmw.WorkspaceAgent(r) + + // Chats are processed by the chat daemon; without it there is + // nothing to refresh. + if api.chatDaemon == nil { + httpapi.Write(ctx, rw, http.StatusOK, agentsdk.RefreshChatContextResponse{}) + return + } + + // Use system context for chat operations since the workspace agent + // scope does not include chat resources. + //nolint:gocritic // Agent needs system access to read/write chat resources. + sysCtx := dbauthz.AsSystemRestricted(ctx) + workspace, err := api.Database.GetWorkspaceByAgentID(sysCtx, workspaceAgent.ID) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to determine workspace from agent token.", + Detail: err.Error(), + }) + return + } + + chats, err := api.Database.GetActiveChatsByAgentID(sysCtx, workspaceAgent.ID) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to list chats for agent.", + Detail: err.Error(), + }) + return + } + + refreshed := 0 + for _, chat := range chats { + // Only re-pin chats owned by this workspace's owner that have + // drifted from the agent's latest snapshot. + if chat.OwnerID != workspace.OwnerID || !chat.ContextDirtySince.Valid { + continue + } + if _, err := api.chatDaemon.RefreshChatContext(sysCtx, chat); err != nil { + api.Logger.Warn(ctx, "failed to refresh chat context for agent", + slog.F("chat_id", chat.ID), + slog.F("agent_id", workspaceAgent.ID), + slog.Error(err), + ) + continue + } + refreshed++ + } + + httpapi.Write(ctx, rw, http.StatusOK, agentsdk.RefreshChatContextResponse{ + Refreshed: refreshed, + }) +} + var ( errNoActiveChats = xerrors.New("no active chats found") errChatNotFound = xerrors.New("chat not found") diff --git a/coderd/x/chatd/context_integration_test.go b/coderd/x/chatd/context_integration_test.go index b6307f6197..145244decd 100644 --- a/coderd/x/chatd/context_integration_test.go +++ b/coderd/x/chatd/context_integration_test.go @@ -269,3 +269,130 @@ func TestChatContextDirtyFromAgentPush(t *testing.T) { require.NotNil(t, got.Context) require.False(t, got.Context.Dirty, "re-push of the pinned hash stays clean") } + +// TestChatContextRefreshFromAgentToken covers the in-workspace +// `coder exp chat context refresh` (no chat argument) path, which authenticates +// with the agent token instead of a user session. The agent endpoint re-pins +// every drifted chat bound to the calling agent to its latest snapshot and +// clears the drift marker, returning how many were refreshed. A chat bound to +// no agent must stay untouched, guarding the agent-scoped query. +func TestChatContextRefreshFromAgentToken(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := coderdtest.NewWithDatabase(t, &coderdtest.Options{ + DeploymentValues: directChatRoutingDeploymentValues(t), + IncludeProvisionerDaemon: true, + }) + user := coderdtest.CreateFirstUser(t, client) + expClient := codersdk.NewExperimentalClient(client) + + // Build a workspace with an agent via the echo provisioner so the agent + // token is accepted by the agent middleware backing the endpoint. + agentToken := uuid.NewString() + version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{ + Parse: echo.ParseComplete, + ProvisionPlan: echo.PlanComplete, + ProvisionApply: echo.ApplyComplete, + ProvisionGraph: echo.ProvisionGraphWithAgent(agentToken), + }) + coderdtest.AwaitTemplateVersionJobCompleted(t, client, version.ID) + template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID) + workspace := coderdtest.CreateWorkspace(t, client, template.ID) + coderdtest.AwaitWorkspaceBuildJobCompleted(t, client, workspace.LatestBuild.ID) + + ws, err := client.Workspace(ctx, workspace.ID) + require.NoError(t, err) + require.Len(t, ws.LatestBuild.Resources, 1) + require.Len(t, ws.LatestBuild.Resources[0].Agents, 1) + agentID := ws.LatestBuild.Resources[0].Agents[0].ID + + // A chat bound to the agent, plus an unrelated chat bound to no agent that + // must stay untouched by the agent-scoped refresh. + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: user.OrganizationID, + OwnerID: user.UserID, + WorkspaceID: uuid.NullUUID{UUID: workspace.ID, Valid: true}, + AgentID: uuid.NullUUID{UUID: agentID, Valid: true}, + LastModelConfigID: model.ID, + Status: database.ChatStatusWaiting, + }) + otherChat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: user.OrganizationID, + OwnerID: user.UserID, + LastModelConfigID: model.ID, + Status: database.ChatStatusWaiting, + }) + + agentsSource := "/home/coder/workspace/AGENTS.md" + instructionResource := func(content string, hash []byte) *agentproto.ContextResource { + return &agentproto.ContextResource{ + Source: agentsSource, + ContentHash: hash, + SizeBytes: uint64(len(content)), + Status: agentproto.ContextResource_OK, + Body: &agentproto.ContextResource_InstructionFile{ + InstructionFile: &agentproto.InstructionFileBody{Content: []byte(content)}, + }, + } + } + + // The agent token drives both the DRPC push and the REST refresh. + agentClient := agentsdk.New(client.URL, agentsdk.WithFixedToken(agentToken)) + aAPI, _, err := agentClient.ConnectRPC210(ctx) + require.NoError(t, err) + defer func() { _ = aAPI.DRPCConn().Close() }() + + // Initial push hydrates the chat to a clean context. + resp, err := aAPI.PushContextState(ctx, &agentproto.PushContextStateRequest{ + Version: 1, + Initial: true, + AggregateHash: []byte{0x01}, + Resources: []*agentproto.ContextResource{instructionResource("hello-v1", []byte{0x11})}, + }) + require.NoError(t, err) + require.True(t, resp.GetAccepted()) + + // With nothing dirty, the agent-token refresh is a no-op. + refresh, err := agentClient.RefreshChatContext(ctx) + require.NoError(t, err) + require.Equal(t, 0, refresh.Refreshed, "no dirty chats to refresh") + + // A second push with a different hash drifts the bound chat dirty. + resp, err = aAPI.PushContextState(ctx, &agentproto.PushContextStateRequest{ + Version: 2, + AggregateHash: []byte{0x02}, + Resources: []*agentproto.ContextResource{instructionResource("hello-v2", []byte{0x22})}, + }) + require.NoError(t, err) + require.True(t, resp.GetAccepted()) + + got, err := expClient.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.NotNil(t, got.Context) + require.True(t, got.Context.Dirty, "second push drifts the chat dirty") + + // The agent-token refresh re-pins every drifted chat bound to the agent. + refresh, err = agentClient.RefreshChatContext(ctx) + require.NoError(t, err) + require.Equal(t, 1, refresh.Refreshed, "the drifted chat is re-pinned") + + got, err = expClient.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.NotNil(t, got.Context) + require.False(t, got.Context.Dirty, "refresh clears the dirty marker") + require.Len(t, got.Context.Resources, 1) + require.Equal(t, agentsSource, got.Context.Resources[0].Source) + + // The agent-less chat is never returned by the agent-scoped query, so it + // must stay unhydrated throughout. + other, err := expClient.GetChat(ctx, otherChat.ID) + require.NoError(t, err) + require.Nil(t, other.Context, "agent-less chat stays untouched") + + // A follow-up refresh with nothing dirty is a no-op again. + refresh, err = agentClient.RefreshChatContext(ctx) + require.NoError(t, err) + require.Equal(t, 0, refresh.Refreshed, "nothing left to refresh") +} diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go index 815f175240..27ca113565 100644 --- a/codersdk/agentsdk/agentsdk.go +++ b/codersdk/agentsdk/agentsdk.go @@ -1021,6 +1021,13 @@ type ClearChatContextResponse struct { ChatID uuid.UUID `json:"chat_id"` } +// RefreshChatContextResponse is the response for refreshing chat context. +type RefreshChatContextResponse struct { + // Refreshed is the number of drifted chats that were re-pinned to the + // agent's latest context snapshot. + Refreshed int `json:"refreshed"` +} + // AddChatContext adds context-file and skill parts to an active chat. func (c *Client) AddChatContext(ctx context.Context, req AddChatContextRequest) (AddChatContextResponse, error) { res, err := c.SDK.Request(ctx, http.MethodPost, "/api/v2/workspaceagents/me/experimental/chat-context", req) @@ -1052,3 +1059,22 @@ func (c *Client) ClearChatContext(ctx context.Context, req ClearChatContextReque var resp ClearChatContextResponse return resp, json.NewDecoder(res.Body).Decode(&resp) } + +// RefreshChatContext re-pins every drifted chat bound to this agent to the +// agent's latest context snapshot, clearing their drift markers. It backs +// the in-workspace `coder exp chat context refresh` (no chat argument), +// which authenticates with the agent token rather than a user session. +func (c *Client) RefreshChatContext(ctx context.Context) (RefreshChatContextResponse, error) { + res, err := c.SDK.Request(ctx, http.MethodPost, "/api/v2/workspaceagents/me/experimental/chat-context/refresh", nil) + if err != nil { + return RefreshChatContextResponse{}, xerrors.Errorf("execute request: %w", err) + } + defer res.Body.Close() + + if res.StatusCode != http.StatusOK { + return RefreshChatContextResponse{}, codersdk.ReadBodyAsError(res) + } + + var resp RefreshChatContextResponse + return resp, json.NewDecoder(res.Body).Decode(&resp) +}