mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
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._
This commit is contained in:
Kyle Carberry
@@ -15122,6 +15122,139 @@ func TestSoftDeleteWorkspaceAgentsByWorkspaceID(t *testing.T) {
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// TestSoftDeleteWorkspaceAgentsPurgesContext verifies that both agent
|
||||
// soft-delete queries hard-delete the agents' pushed context rows
|
||||
// (workspace_agent_context_snapshots and
|
||||
// workspace_agent_context_resources). Agents are only ever
|
||||
// soft-deleted, so without this the context rows would accumulate
|
||||
// forever.
|
||||
func TestSoftDeleteWorkspaceAgentsPurgesContext(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
db, _ := dbtestutil.NewDB(t)
|
||||
ctx := testutil.Context(t, testutil.WaitShort)
|
||||
|
||||
user := dbgen.User(t, db, database.User{})
|
||||
org := dbgen.Organization(t, db, database.Organization{})
|
||||
tpl := dbgen.Template(t, db, database.Template{
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
})
|
||||
tplVersion := dbgen.TemplateVersion(t, db, database.TemplateVersion{
|
||||
TemplateID: uuid.NullUUID{UUID: tpl.ID, Valid: true},
|
||||
OrganizationID: org.ID,
|
||||
CreatedBy: user.ID,
|
||||
})
|
||||
|
||||
type buildBundle struct {
|
||||
buildID uuid.UUID
|
||||
agentID uuid.UUID
|
||||
agent database.WorkspaceAgent
|
||||
}
|
||||
|
||||
newBuild := func(t *testing.T, wsID uuid.UUID, buildNumber int32) buildBundle {
|
||||
t.Helper()
|
||||
job := dbgen.ProvisionerJob(t, db, nil, database.ProvisionerJob{
|
||||
OrganizationID: org.ID,
|
||||
Type: database.ProvisionerJobTypeWorkspaceBuild,
|
||||
})
|
||||
build := dbgen.WorkspaceBuild(t, db, database.WorkspaceBuild{
|
||||
WorkspaceID: wsID,
|
||||
JobID: job.ID,
|
||||
TemplateVersionID: tplVersion.ID,
|
||||
BuildNumber: buildNumber,
|
||||
Transition: database.WorkspaceTransitionStart,
|
||||
})
|
||||
resource := dbgen.WorkspaceResource(t, db, database.WorkspaceResource{JobID: job.ID})
|
||||
agent := dbgen.WorkspaceAgent(t, db, database.WorkspaceAgent{ResourceID: resource.ID})
|
||||
return buildBundle{buildID: build.ID, agentID: agent.ID, agent: agent}
|
||||
}
|
||||
|
||||
pushContext := func(t *testing.T, agentID uuid.UUID) {
|
||||
t.Helper()
|
||||
_, err := db.UpsertWorkspaceAgentContextSnapshot(ctx, database.UpsertWorkspaceAgentContextSnapshotParams{
|
||||
WorkspaceAgentID: agentID,
|
||||
Version: 1,
|
||||
AggregateHash: []byte{0x01},
|
||||
ReceivedAt: dbtime.Now(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
_, err = db.UpsertWorkspaceAgentContextResource(ctx, database.UpsertWorkspaceAgentContextResourceParams{
|
||||
WorkspaceAgentID: agentID,
|
||||
Source: "/workspace/AGENTS.md",
|
||||
BodyKind: database.WorkspaceAgentContextBodyKindInstructionFile,
|
||||
Body: []byte(`{}`),
|
||||
ContentHash: []byte{0x02},
|
||||
SizeBytes: 2,
|
||||
Status: database.WorkspaceAgentContextResourceStatusOk,
|
||||
Now: dbtime.Now(),
|
||||
})
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
hasContext := func(t *testing.T, agentID uuid.UUID) bool {
|
||||
t.Helper()
|
||||
_, err := db.GetLatestWorkspaceAgentContextSnapshot(ctx, agentID)
|
||||
if errors.Is(err, sql.ErrNoRows) {
|
||||
resources, err := db.ListWorkspaceAgentContextResources(ctx, agentID)
|
||||
require.NoError(t, err)
|
||||
require.Empty(t, resources, "snapshot and resource rows must be deleted together")
|
||||
return false
|
||||
}
|
||||
require.NoError(t, err)
|
||||
return true
|
||||
}
|
||||
|
||||
wsA := dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: tpl.ID,
|
||||
OwnerID: user.ID,
|
||||
}).ID
|
||||
wsB := dbgen.Workspace(t, db, database.WorkspaceTable{
|
||||
OrganizationID: org.ID,
|
||||
TemplateID: tpl.ID,
|
||||
OwnerID: user.ID,
|
||||
}).ID
|
||||
|
||||
a1 := newBuild(t, wsA, 1)
|
||||
a2 := newBuild(t, wsA, 2)
|
||||
b1 := newBuild(t, wsB, 1)
|
||||
|
||||
pushContext(t, a1.agentID)
|
||||
pushContext(t, a2.agentID)
|
||||
pushContext(t, b1.agentID)
|
||||
|
||||
// Soft-deleting wsA's prior agents purges a1's context but leaves
|
||||
// the current build's agent and other workspaces untouched.
|
||||
err := db.SoftDeletePriorWorkspaceAgents(ctx, database.SoftDeletePriorWorkspaceAgentsParams{
|
||||
WorkspaceID: wsA,
|
||||
CurrentBuildID: a2.buildID,
|
||||
})
|
||||
require.NoError(t, err)
|
||||
assert.False(t, hasContext(t, a1.agentID), "prior build agent context must be purged")
|
||||
assert.True(t, hasContext(t, a2.agentID), "current build agent context must remain")
|
||||
assert.True(t, hasContext(t, b1.agentID), "other workspace agent context must remain")
|
||||
|
||||
// Soft-deleting all of wsB's agents purges b1's context.
|
||||
err = db.SoftDeleteWorkspaceAgentsByWorkspaceID(ctx, wsB)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, hasContext(t, a2.agentID), "other workspace agent context must remain")
|
||||
assert.False(t, hasContext(t, b1.agentID), "deleted workspace agent context must be purged")
|
||||
|
||||
// Removing a sub-agent mid-build via DeleteWorkspaceSubAgentByID purges
|
||||
// only that sub-agent's context. The rebuild-time queries skip
|
||||
// already-deleted agents, so this is the sole cleanup opportunity.
|
||||
c1 := newBuild(t, wsA, 3)
|
||||
subAgent := dbgen.WorkspaceSubAgent(t, db, c1.agent, database.WorkspaceAgent{})
|
||||
pushContext(t, c1.agentID)
|
||||
pushContext(t, subAgent.ID)
|
||||
|
||||
err = db.DeleteWorkspaceSubAgentByID(ctx, subAgent.ID)
|
||||
require.NoError(t, err)
|
||||
assert.True(t, hasContext(t, c1.agentID), "parent agent context must remain")
|
||||
assert.False(t, hasContext(t, subAgent.ID), "deleted sub-agent context must be purged")
|
||||
}
|
||||
|
||||
func TestAIGatewayKeysTableConstraints(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user