mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat(coderd): add task log snapshot storage endpoint (#21644)
This change adds a POST /workspaceagents/me/tasks/{task}/log-snapshot
endpoint for agents to upload task conversation history during
workspace shutdown. This allows users to view task logs even when the
workspace is stopped.
The endpoint accepts agentapi format payloads (typically last 10
messages, max 64KB), wraps them in a format envelope, and upserts to the
task_snapshots table. Uses agent token auth and validates the task
belongs to the agent's workspace.
Closes coder/internal#1253
This commit is contained in:
@@ -3055,6 +3055,25 @@ func (q *querier) GetTaskByWorkspaceID(ctx context.Context, workspaceID uuid.UUI
|
||||
return fetch(q.log, q.auth, q.db.GetTaskByWorkspaceID)(ctx, workspaceID)
|
||||
}
|
||||
|
||||
func (q *querier) GetTaskSnapshot(ctx context.Context, taskID uuid.UUID) (database.TaskSnapshot, error) {
|
||||
// Fetch task to build RBAC object for authorization.
|
||||
task, err := q.GetTaskByID(ctx, taskID)
|
||||
if err != nil {
|
||||
return database.TaskSnapshot{}, err
|
||||
}
|
||||
|
||||
obj := rbac.ResourceTask.
|
||||
WithID(task.ID).
|
||||
WithOwner(task.OwnerID.String()).
|
||||
InOrg(task.OrganizationID)
|
||||
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, obj); err != nil {
|
||||
return database.TaskSnapshot{}, err
|
||||
}
|
||||
|
||||
return q.db.GetTaskSnapshot(ctx, taskID)
|
||||
}
|
||||
|
||||
func (q *querier) GetTelemetryItem(ctx context.Context, key string) (database.TelemetryItem, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
|
||||
return database.TelemetryItem{}, err
|
||||
@@ -6024,6 +6043,25 @@ func (q *querier) UpsertTailnetTunnel(ctx context.Context, arg database.UpsertTa
|
||||
return q.db.UpsertTailnetTunnel(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertTaskSnapshot(ctx context.Context, arg database.UpsertTaskSnapshotParams) error {
|
||||
// Fetch task to build RBAC object for authorization.
|
||||
task, err := q.GetTaskByID(ctx, arg.TaskID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
obj := rbac.ResourceTask.
|
||||
WithID(task.ID).
|
||||
WithOwner(task.OwnerID.String()).
|
||||
InOrg(task.OrganizationID)
|
||||
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, obj); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return q.db.UpsertTaskSnapshot(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertTaskWorkspaceApp(ctx context.Context, arg database.UpsertTaskWorkspaceAppParams) (database.TaskWorkspaceApp, error) {
|
||||
// Fetch the task to derive the RBAC object and authorize update on it.
|
||||
task, err := q.db.GetTaskByID(ctx, arg.TaskID)
|
||||
|
||||
@@ -2547,6 +2547,24 @@ func (s *MethodTestSuite) TestTasks() {
|
||||
dbm.EXPECT().ListTasks(gomock.Any(), gomock.Any()).Return([]database.Task{t1, t2}, nil).AnyTimes()
|
||||
check.Args(database.ListTasksParams{}).Asserts(t1, policy.ActionRead, t2, policy.ActionRead).Returns([]database.Task{t1, t2})
|
||||
}))
|
||||
s.Run("GetTaskSnapshot", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
task := testutil.Fake(s.T(), faker, database.Task{})
|
||||
snapshot := testutil.Fake(s.T(), faker, database.TaskSnapshot{TaskID: task.ID})
|
||||
dbm.EXPECT().GetTaskByID(gomock.Any(), task.ID).Return(task, nil).AnyTimes()
|
||||
dbm.EXPECT().GetTaskSnapshot(gomock.Any(), task.ID).Return(snapshot, nil).AnyTimes()
|
||||
check.Args(task.ID).Asserts(task, policy.ActionRead, task, policy.ActionRead).Returns(snapshot)
|
||||
}))
|
||||
s.Run("UpsertTaskSnapshot", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
task := testutil.Fake(s.T(), faker, database.Task{})
|
||||
arg := database.UpsertTaskSnapshotParams{
|
||||
TaskID: task.ID,
|
||||
LogSnapshot: []byte(`{"format":"agentapi","data":[]}`),
|
||||
LogSnapshotCreatedAt: dbtime.Now(),
|
||||
}
|
||||
dbm.EXPECT().GetTaskByID(gomock.Any(), task.ID).Return(task, nil).AnyTimes()
|
||||
dbm.EXPECT().UpsertTaskSnapshot(gomock.Any(), arg).Return(nil).AnyTimes()
|
||||
check.Args(arg).Asserts(task, policy.ActionRead, task, policy.ActionUpdate).Returns()
|
||||
}))
|
||||
}
|
||||
|
||||
func (s *MethodTestSuite) TestProvisionerKeys() {
|
||||
|
||||
@@ -1758,6 +1758,14 @@ func (m queryMetricsStore) GetTaskByWorkspaceID(ctx context.Context, workspaceID
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetTaskSnapshot(ctx context.Context, taskID uuid.UUID) (database.TaskSnapshot, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetTaskSnapshot(ctx, taskID)
|
||||
m.queryLatencies.WithLabelValues("GetTaskSnapshot").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetTaskSnapshot").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetTelemetryItem(ctx context.Context, key string) (database.TelemetryItem, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetTelemetryItem(ctx, key)
|
||||
@@ -4189,6 +4197,14 @@ func (m queryMetricsStore) UpsertTailnetTunnel(ctx context.Context, arg database
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpsertTaskSnapshot(ctx context.Context, arg database.UpsertTaskSnapshotParams) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.UpsertTaskSnapshot(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("UpsertTaskSnapshot").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertTaskSnapshot").Inc()
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpsertTaskWorkspaceApp(ctx context.Context, arg database.UpsertTaskWorkspaceAppParams) (database.TaskWorkspaceApp, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.UpsertTaskWorkspaceApp(ctx, arg)
|
||||
|
||||
@@ -3254,6 +3254,21 @@ func (mr *MockStoreMockRecorder) GetTaskByWorkspaceID(ctx, workspaceID any) *gom
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTaskByWorkspaceID", reflect.TypeOf((*MockStore)(nil).GetTaskByWorkspaceID), ctx, workspaceID)
|
||||
}
|
||||
|
||||
// GetTaskSnapshot mocks base method.
|
||||
func (m *MockStore) GetTaskSnapshot(ctx context.Context, taskID uuid.UUID) (database.TaskSnapshot, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetTaskSnapshot", ctx, taskID)
|
||||
ret0, _ := ret[0].(database.TaskSnapshot)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetTaskSnapshot indicates an expected call of GetTaskSnapshot.
|
||||
func (mr *MockStoreMockRecorder) GetTaskSnapshot(ctx, taskID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetTaskSnapshot", reflect.TypeOf((*MockStore)(nil).GetTaskSnapshot), ctx, taskID)
|
||||
}
|
||||
|
||||
// GetTelemetryItem mocks base method.
|
||||
func (m *MockStore) GetTelemetryItem(ctx context.Context, key string) (database.TelemetryItem, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -7817,6 +7832,20 @@ func (mr *MockStoreMockRecorder) UpsertTailnetTunnel(ctx, arg any) *gomock.Call
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertTailnetTunnel", reflect.TypeOf((*MockStore)(nil).UpsertTailnetTunnel), ctx, arg)
|
||||
}
|
||||
|
||||
// UpsertTaskSnapshot mocks base method.
|
||||
func (m *MockStore) UpsertTaskSnapshot(ctx context.Context, arg database.UpsertTaskSnapshotParams) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "UpsertTaskSnapshot", ctx, arg)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// UpsertTaskSnapshot indicates an expected call of UpsertTaskSnapshot.
|
||||
func (mr *MockStoreMockRecorder) UpsertTaskSnapshot(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertTaskSnapshot", reflect.TypeOf((*MockStore)(nil).UpsertTaskSnapshot), ctx, arg)
|
||||
}
|
||||
|
||||
// UpsertTaskWorkspaceApp mocks base method.
|
||||
func (m *MockStore) UpsertTaskWorkspaceApp(ctx context.Context, arg database.UpsertTaskWorkspaceAppParams) (database.TaskWorkspaceApp, error) {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
@@ -353,6 +353,7 @@ type sqlcQuerier interface {
|
||||
GetTaskByID(ctx context.Context, id uuid.UUID) (Task, error)
|
||||
GetTaskByOwnerIDAndName(ctx context.Context, arg GetTaskByOwnerIDAndNameParams) (Task, error)
|
||||
GetTaskByWorkspaceID(ctx context.Context, workspaceID uuid.UUID) (Task, error)
|
||||
GetTaskSnapshot(ctx context.Context, taskID uuid.UUID) (TaskSnapshot, error)
|
||||
GetTelemetryItem(ctx context.Context, key string) (TelemetryItem, error)
|
||||
GetTelemetryItems(ctx context.Context) ([]TelemetryItem, error)
|
||||
// GetTemplateAppInsights returns the aggregate usage of each app in a given
|
||||
@@ -772,6 +773,7 @@ type sqlcQuerier interface {
|
||||
UpsertTailnetCoordinator(ctx context.Context, id uuid.UUID) (TailnetCoordinator, error)
|
||||
UpsertTailnetPeer(ctx context.Context, arg UpsertTailnetPeerParams) (TailnetPeer, error)
|
||||
UpsertTailnetTunnel(ctx context.Context, arg UpsertTailnetTunnelParams) (TailnetTunnel, error)
|
||||
UpsertTaskSnapshot(ctx context.Context, arg UpsertTaskSnapshotParams) error
|
||||
UpsertTaskWorkspaceApp(ctx context.Context, arg UpsertTaskWorkspaceAppParams) (TaskWorkspaceApp, error)
|
||||
UpsertTelemetryItem(ctx context.Context, arg UpsertTelemetryItemParams) error
|
||||
// This query aggregates the workspace_agent_stats and workspace_app_stats data
|
||||
|
||||
@@ -13183,6 +13183,22 @@ func (q *sqlQuerier) GetTaskByWorkspaceID(ctx context.Context, workspaceID uuid.
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getTaskSnapshot = `-- name: GetTaskSnapshot :one
|
||||
SELECT
|
||||
task_id, log_snapshot, log_snapshot_created_at
|
||||
FROM
|
||||
task_snapshots
|
||||
WHERE
|
||||
task_id = $1
|
||||
`
|
||||
|
||||
func (q *sqlQuerier) GetTaskSnapshot(ctx context.Context, taskID uuid.UUID) (TaskSnapshot, error) {
|
||||
row := q.db.QueryRowContext(ctx, getTaskSnapshot, taskID)
|
||||
var i TaskSnapshot
|
||||
err := row.Scan(&i.TaskID, &i.LogSnapshot, &i.LogSnapshotCreatedAt)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const insertTask = `-- name: InsertTask :one
|
||||
INSERT INTO tasks
|
||||
(id, organization_id, owner_id, name, display_name, workspace_id, template_version_id, template_parameters, prompt, created_at)
|
||||
@@ -13373,6 +13389,29 @@ func (q *sqlQuerier) UpdateTaskWorkspaceID(ctx context.Context, arg UpdateTaskWo
|
||||
return i, err
|
||||
}
|
||||
|
||||
const upsertTaskSnapshot = `-- name: UpsertTaskSnapshot :exec
|
||||
INSERT INTO
|
||||
task_snapshots (task_id, log_snapshot, log_snapshot_created_at)
|
||||
VALUES
|
||||
($1, $2, $3)
|
||||
ON CONFLICT
|
||||
(task_id)
|
||||
DO UPDATE SET
|
||||
log_snapshot = EXCLUDED.log_snapshot,
|
||||
log_snapshot_created_at = EXCLUDED.log_snapshot_created_at
|
||||
`
|
||||
|
||||
type UpsertTaskSnapshotParams struct {
|
||||
TaskID uuid.UUID `db:"task_id" json:"task_id"`
|
||||
LogSnapshot json.RawMessage `db:"log_snapshot" json:"log_snapshot"`
|
||||
LogSnapshotCreatedAt time.Time `db:"log_snapshot_created_at" json:"log_snapshot_created_at"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) UpsertTaskSnapshot(ctx context.Context, arg UpsertTaskSnapshotParams) error {
|
||||
_, err := q.db.ExecContext(ctx, upsertTaskSnapshot, arg.TaskID, arg.LogSnapshot, arg.LogSnapshotCreatedAt)
|
||||
return err
|
||||
}
|
||||
|
||||
const upsertTaskWorkspaceApp = `-- name: UpsertTaskWorkspaceApp :one
|
||||
INSERT INTO task_workspace_apps
|
||||
(task_id, workspace_build_number, workspace_agent_id, workspace_app_id)
|
||||
|
||||
@@ -75,3 +75,22 @@ WHERE
|
||||
id = @id::uuid
|
||||
AND deleted_at IS NULL
|
||||
RETURNING *;
|
||||
|
||||
-- name: UpsertTaskSnapshot :exec
|
||||
INSERT INTO
|
||||
task_snapshots (task_id, log_snapshot, log_snapshot_created_at)
|
||||
VALUES
|
||||
($1, $2, $3)
|
||||
ON CONFLICT
|
||||
(task_id)
|
||||
DO UPDATE SET
|
||||
log_snapshot = EXCLUDED.log_snapshot,
|
||||
log_snapshot_created_at = EXCLUDED.log_snapshot_created_at;
|
||||
|
||||
-- name: GetTaskSnapshot :one
|
||||
SELECT
|
||||
*
|
||||
FROM
|
||||
task_snapshots
|
||||
WHERE
|
||||
task_id = $1;
|
||||
|
||||
Reference in New Issue
Block a user