Files
coder/agent/agentcontext/api_test.go
T
Kyle Carberry b439b06ee6 feat: persist agent-pushed workspace context snapshots in coderd (#26145)
Replaces the v2.10 `PushContextState` stub with a real coderd write
path. Phase 1 of the chat-side persistence story; nothing reads these
rows yet.

Follows [#25983](https://github.com/coder/coder/pull/25983) and unblocks
[CODAGT-569](https://linear.app/codercom/issue/CODAGT-569/enable-agent-api-v210-pushcontextstate-bump-currentminor-wire-coderd).

## What ships

### Schema (`000517_workspace_agent_context.{up,down}.sql`)

Two new tables plus `api_key_scope` enum extensions:

- `workspace_agent_context_snapshots` (PK `workspace_agent_id` to
`workspace_agents(id) ON DELETE CASCADE`): one row per agent,
overwritten per push. Holds `version`, `schema_version`,
`aggregate_hash`, `snapshot_error`, `received_at`.
- `workspace_agent_context_resources` (PK `(workspace_agent_id,
source)`): per-resource state. `body_kind` and `status` are `TEXT` +
`CHECK` so adding new wire kinds (the RFC's reserved
PLUGIN/HOOK/SUBAGENT/COMMAND) is a one-line CHECK update plus a Go
switch case.

### SQLC queries (`coderd/database/queries/workspaceagentcontext.sql`)

- `UpsertWorkspaceAgentContextSnapshot`
- `UpsertWorkspaceAgentContextResource`
- `DeleteStaleWorkspaceAgentContextResources`
(delete-where-source-not-in)
- `GetLatestWorkspaceAgentContextSnapshot`
- `ListWorkspaceAgentContextResources`

### Handler (`coderd/agentapi/context.go`)

`ContextAPI` is a new sub-API. `PushContextState`:

1. Rejects `schema_version > 1` with a non-`Unimplemented` error so a
forward-incompatible agent fails loudly during rollout instead of
slipping into the permanent fallback path the `Unimplemented`
translation reserves for old coderd deployments.
2. Validates resources: no empty/duplicate sources, every variant maps
to a known body kind, every status maps to a known enum value, the
`Body` oneof is set (even when status is non-OK, mirroring the wire
guarantee so coderd can attribute failures to a known kind).
3. Inside `Database.InTx`, reads the existing snapshot. If the push is
not `initial` and `version` is not strictly greater, returns `accepted =
false` and leaves stored state untouched. Otherwise upserts the snapshot
row, upserts each resource, then runs the stale-source prune so the
snapshot and resource rows always agree.
4. Returns `accepted = true` on success.

Resource bodies are stored as `protojson(body oneof variant)` in `body
JSONB` with `body_kind` as the discriminator. Adding a new field to an
existing variant is zero work since `protojson` tolerates new fields;
adding a new variant is a CHECK + switch case.

### RBAC + dbauthz

- New `ResourceWorkspaceAgentContext` (Create/Read/Update/Delete).
- New `SubjectTypeAgentContext` plus `subjectAgentContext` system role
and `dbauthz.AsAgentContext` helper. The push handler elevates to this
subject; the agent's own role does not get direct write access to the
table.
- New `workspace_agent_context:*` API key scopes registered in the enum
migration; internal-only (not added to `externalLowLevel`).

### Audit

These rows are agent-pushed state, not user-authored. They are
intentionally not added to `AuditActionMap` and not enumerated in
`enterprise/audit/table.go`, matching `boundary_logs`,
`workspace_agent_memory_resource_monitor`, etc. `enterprise/audit` tests
pass unchanged.

## Tests

- `coderd/agentapi/context_test.go`: 12 subtests covering
accepts/rejects (schema version, empty/duplicate source, unknown status,
missing body), version semantics (stale dropped, same-version replay
dropped, `initial=true` overwrites lower version), variant coverage,
non-OK status persistence, and the empty-active-set prune case.
- `coderd/database/dbauthz/dbauthz_test.go`: 5 `MethodTestSuite` cases
covering the new queries.
- `coderd/rbac/roles_test.go`: `WorkspaceAgentContext` permission row
asserting no human role currently has access.
-
`coderd/database/migrations/testdata/fixtures/000517_workspace_agent_context.up.sql`:
one snapshot + one resource per known body kind plus a non-OK status, so
the migration test suite never lands with these tables empty.

## Out of scope (later phases)

- Chat hydration (`chats.context_aggregate_hash`,
`last_injected_context`).
- Dirty-bit fan-out and `PUT /chats/{id}/context`.
- Agent-side `POST /api/v0/context/resync` barrier and the `coder exp
chat context` CLI.
- `codersdk` chat-context wire types and the dashboard Sources drawer.
- Removal of the chatd per-turn pull fallback.

## Compat property

This is a pure write path. If anything here returns errors the agent's
`RunPush` loop backs off, no chat behavior changes, and the workspace
keeps behaving exactly like it did before v2.10.

<details>
<summary>Implementation plan and decision log</summary>

Key design calls:

1. **Concurrency**: Accept iff `req.Initial || req.Version >
existing.Version`. The strict RFC reading ("version comparison is
authoritative") locks restarted agents out because their per-process
counter resets to 1; honoring `initial=true` reflects the real reboot
reality while still rejecting steady-state replays/out-of-order pushes.
2. **Body encoding**: `protojson` over the oneof variant body proto,
stored in JSONB with `body_kind` discriminator. Structured at the API/Go
layer, schema-tolerant at the storage layer, and Phase 2 readers
round-trip back via `protojson.Unmarshal`.
3. **Schema version rejection**: returns a normal error, not
`Unimplemented`. The agent's `RunPush` loop only short-circuits on
`Unimplemented`; that escape hatch is reserved for old coderd
deployments. A forward-incompatible agent should retry-and-back-off, not
flip the connection into permanent fallback.
4. **Validation strictness**: empty sources, duplicate sources,
`STATUS_UNSPECIFIED`, and missing `Body` oneof variants are rejected
before any write so a misbehaving agent cannot poison the snapshot
table. Phase 2 readers can trust every row maps to a known proto
variant.

</details>

_This PR was authored by Coder Agents on Kyle Carberry's behalf._
2026-06-15 09:38:52 -07:00

177 lines
5.3 KiB
Go

package agentcontext_test
import (
"bytes"
"encoding/json"
"io"
"net/http"
"net/http/httptest"
"net/url"
"path/filepath"
"testing"
"github.com/stretchr/testify/require"
"github.com/coder/coder/v2/agent/agentcontext"
"github.com/coder/coder/v2/testutil"
)
func newAPITestServer(t *testing.T, opts agentcontext.ManagerOptions) (*httptest.Server, *agentcontext.Manager) {
t.Helper()
m := newTestManager(t, opts)
api := agentcontext.NewAPI(m)
srv := httptest.NewServer(api.Routes())
t.Cleanup(srv.Close)
return srv, m
}
// doRequest issues an HTTP request bounded by testutil.WaitShort
// and returns the status code and response body. The response
// body is closed before doRequest returns.
func doRequest(t *testing.T, method, requrl string, body io.Reader) (int, []byte) {
t.Helper()
ctx := testutil.Context(t, testutil.WaitShort)
req, err := http.NewRequestWithContext(ctx, method, requrl, body)
require.NoError(t, err)
if body != nil {
req.Header.Set("Content-Type", "application/json")
}
res, err := http.DefaultClient.Do(req) //nolint:bodyclose // closed below.
require.NoError(t, err)
defer res.Body.Close()
bodyBytes, err := io.ReadAll(res.Body)
require.NoError(t, err)
return res.StatusCode, bodyBytes
}
func TestAPI_ListSourcesEmpty(t *testing.T) {
t.Parallel()
srv, _ := newAPITestServer(t, agentcontext.ManagerOptions{
WorkingDir: func() string { return t.TempDir() },
})
status, body := doRequest(t, http.MethodGet, srv.URL+"/sources", nil)
require.Equal(t, http.StatusOK, status)
var got []agentcontext.SourceResponse
require.NoError(t, json.Unmarshal(body, &got))
require.Empty(t, got)
}
func TestAPI_AddAndListSource(t *testing.T) {
t.Parallel()
wd := t.TempDir()
src := testutil.TempDirResolved(t)
srv, _ := newAPITestServer(t, agentcontext.ManagerOptions{
WorkingDir: func() string { return wd },
AllowedRoots: []string{wd, src},
})
body, _ := json.Marshal(agentcontext.SourceRequest{Path: src})
status, addBody := doRequest(t, http.MethodPost, srv.URL+"/sources", bytes.NewReader(body))
require.Equal(t, http.StatusCreated, status)
var created agentcontext.SourceResponse
require.NoError(t, json.Unmarshal(addBody, &created))
require.Equal(t, src, created.Path)
// List should show the new source.
listStatus, listBody := doRequest(t, http.MethodGet, srv.URL+"/sources", nil)
require.Equal(t, http.StatusOK, listStatus)
var list []agentcontext.SourceResponse
require.NoError(t, json.Unmarshal(listBody, &list))
require.Len(t, list, 1)
require.Equal(t, src, list[0].Path)
}
func TestAPI_AddSourceRejected(t *testing.T) {
t.Parallel()
wd := t.TempDir()
outside := t.TempDir()
srv, _ := newAPITestServer(t, agentcontext.ManagerOptions{
WorkingDir: func() string { return wd },
AllowedRoots: []string{wd},
})
body, _ := json.Marshal(agentcontext.SourceRequest{Path: outside})
status, _ := doRequest(t, http.MethodPost, srv.URL+"/sources", bytes.NewReader(body))
require.Equal(t, http.StatusBadRequest, status)
}
func TestAPI_GetAndDeleteSource(t *testing.T) {
t.Parallel()
wd := t.TempDir()
src := testutil.TempDirResolved(t)
srv, m := newAPITestServer(t, agentcontext.ManagerOptions{
WorkingDir: func() string { return wd },
AllowedRoots: []string{wd, src},
})
_, err := m.AddSource(agentcontext.Source{Path: src})
require.NoError(t, err)
status, body := doRequest(t, http.MethodGet, srv.URL+"/sources/"+url.PathEscape(src), nil)
require.Equal(t, http.StatusOK, status)
var got agentcontext.SourceResponse
require.NoError(t, json.Unmarshal(body, &got))
require.Equal(t, src, got.Path)
delStatus, _ := doRequest(t, http.MethodDelete, srv.URL+"/sources/"+url.PathEscape(src), nil)
require.Equal(t, http.StatusNoContent, delStatus)
require.Empty(t, m.Sources())
}
func TestAPI_GetSourceNotFound(t *testing.T) {
t.Parallel()
srv, _ := newAPITestServer(t, agentcontext.ManagerOptions{
WorkingDir: func() string { return t.TempDir() },
})
status, _ := doRequest(t, http.MethodGet, srv.URL+"/sources/"+url.PathEscape("/never-added"), nil)
require.Equal(t, http.StatusNotFound, status)
}
func TestAPI_DeleteSourceNotFound(t *testing.T) {
t.Parallel()
srv, _ := newAPITestServer(t, agentcontext.ManagerOptions{
WorkingDir: func() string { return t.TempDir() },
})
status, _ := doRequest(t, http.MethodDelete, srv.URL+"/sources/"+url.PathEscape("/never-added"), nil)
require.Equal(t, http.StatusNotFound, status)
}
func TestAPI_Resync(t *testing.T) {
t.Parallel()
wd := t.TempDir()
mustWriteFile(t, filepath.Join(wd, "AGENTS.md"), "hello")
srv, _ := newAPITestServer(t, agentcontext.ManagerOptions{
WorkingDir: func() string { return wd },
})
status, body := doRequest(t, http.MethodPost, srv.URL+"/resync", nil)
require.Equal(t, http.StatusOK, status)
var snap agentcontext.SnapshotResponse
require.NoError(t, json.Unmarshal(body, &snap))
require.NotEmpty(t, snap.AggregateHash)
require.Len(t, snap.Resources, 1)
require.Equal(t, "instruction_file", snap.Resources[0].Kind)
require.Equal(t, "ok", snap.Resources[0].Status)
}
func TestAPI_AddSourceMalformedBody(t *testing.T) {
t.Parallel()
srv, _ := newAPITestServer(t, agentcontext.ManagerOptions{
WorkingDir: func() string { return t.TempDir() },
})
status, _ := doRequest(t, http.MethodPost, srv.URL+"/sources", bytes.NewReader([]byte("{not json")))
require.Equal(t, http.StatusBadRequest, status)
}