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
2026-09-24 15:04:27 +08:00 · 2026-01-27 11:09:24 +02:00
parent f2e998848e
commit 25d7f27cdb
13 changed files with 745 additions and 8 deletions
@@ -3,6 +3,7 @@ package coderd
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net"
@@ -12,10 +13,12 @@ import (
 	"strings"
 	"time"

+	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
 	"golang.org/x/xerrors"

-	aiagentapi "github.com/coder/agentapi-sdk-go"
+	"cdr.dev/slog/v3"
+	agentapisdk "github.com/coder/agentapi-sdk-go"
 	"github.com/coder/coder/v2/coderd/audit"
 	"github.com/coder/coder/v2/coderd/database"
 	"github.com/coder/coder/v2/coderd/database/dbtime"
@@ -740,7 +743,7 @@ func (api *API) taskSend(rw http.ResponseWriter, r *http.Request) {
 	}

 	if err := api.authAndDoWithTaskAppClient(r, task, func(ctx context.Context, client *http.Client, appURL *url.URL) error {
-		agentAPIClient, err := aiagentapi.NewClient(appURL.String(), aiagentapi.WithHTTPClient(client))
+		agentAPIClient, err := agentapisdk.NewClient(appURL.String(), agentapisdk.WithHTTPClient(client))
 		if err != nil {
 			return httperror.NewResponseError(http.StatusBadGateway, codersdk.Response{
 				Message: "Failed to create agentapi client.",
@@ -756,16 +759,16 @@ func (api *API) taskSend(rw http.ResponseWriter, r *http.Request) {
 			})
 		}

-		if statusResp.Status != aiagentapi.StatusStable {
+		if statusResp.Status != agentapisdk.StatusStable {
 			return httperror.NewResponseError(http.StatusBadGateway, codersdk.Response{
 				Message: "Task app is not ready to accept input.",
 				Detail:  fmt.Sprintf("Status: %s", statusResp.Status),
 			})
 		}

-		_, err = agentAPIClient.PostMessage(ctx, aiagentapi.PostMessageParams{
+		_, err = agentAPIClient.PostMessage(ctx, agentapisdk.PostMessageParams{
 			Content: req.Input,
-			Type:    aiagentapi.MessageTypeUser,
+			Type:    agentapisdk.MessageTypeUser,
 		})
 		if err != nil {
 			return httperror.NewResponseError(http.StatusBadGateway, codersdk.Response{
@@ -798,7 +801,7 @@ func (api *API) taskLogs(rw http.ResponseWriter, r *http.Request) {

 	var out codersdk.TaskLogsResponse
 	if err := api.authAndDoWithTaskAppClient(r, task, func(ctx context.Context, client *http.Client, appURL *url.URL) error {
-		agentAPIClient, err := aiagentapi.NewClient(appURL.String(), aiagentapi.WithHTTPClient(client))
+		agentAPIClient, err := agentapisdk.NewClient(appURL.String(), agentapisdk.WithHTTPClient(client))
 		if err != nil {
 			return httperror.NewResponseError(http.StatusBadGateway, codersdk.Response{
 				Message: "Failed to create agentapi client.",
@@ -818,9 +821,9 @@ func (api *API) taskLogs(rw http.ResponseWriter, r *http.Request) {
 		for _, m := range messagesResp.Messages {
 			var typ codersdk.TaskLogType
 			switch m.Role {
-			case aiagentapi.RoleUser:
+			case agentapisdk.RoleUser:
 				typ = codersdk.TaskLogTypeInput
-			case aiagentapi.RoleAgent:
+			case agentapisdk.RoleAgent:
 				typ = codersdk.TaskLogTypeOutput
 			default:
 				return httperror.NewResponseError(http.StatusBadGateway, codersdk.Response{
@@ -950,3 +953,165 @@ func (api *API) authAndDoWithTaskAppClient(
 	}
 	return do(ctx, client, parsedURL)
 }
+
+const (
+	// taskSnapshotMaxSize is the maximum size for task log snapshots (64KB).
+	// Protects against excessive memory usage and database payload sizes.
+	taskSnapshotMaxSize = 64 * 1024
+)
+
+// TaskLogSnapshotEnvelope wraps a task log snapshot with format metadata.
+type TaskLogSnapshotEnvelope struct {
+	Format string `json:"format"`
+	Data   any    `json:"data"`
+}
+
+// @Summary Upload task log snapshot
+// @ID upload-task-log-snapshot
+// @Security CoderSessionToken
+// @Accept json
+// @Tags Tasks
+// @Param task path string true "Task ID" format(uuid)
+// @Param format query string true "Snapshot format" enums(agentapi)
+// @Param request body object true "Raw snapshot payload (structure depends on format parameter)"
+// @Success 204
+// @Router /workspaceagents/me/tasks/{task}/log-snapshot [post]
+func (api *API) postWorkspaceAgentTaskLogSnapshot(rw http.ResponseWriter, r *http.Request) {
+	var (
+		ctx         = r.Context()
+		latestBuild = httpmw.LatestBuild(r)
+	)
+
+	// Parse task ID from path.
+	taskIDStr := chi.URLParam(r, "task")
+	taskID, err := uuid.Parse(taskIDStr)
+	if err != nil {
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Invalid task ID format.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	// Validate format parameter (required).
+	p := httpapi.NewQueryParamParser().RequiredNotEmpty("format")
+	format := p.String(r.URL.Query(), "", "format")
+	p.ErrorExcessParams(r.URL.Query())
+	if len(p.Errors) > 0 {
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message:     "Invalid query parameters.",
+			Validations: p.Errors,
+		})
+		return
+	}
+	if format != "agentapi" {
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Invalid format parameter.",
+			Detail:  fmt.Sprintf(`Only "agentapi" format is currently supported, got %q.`, format),
+		})
+		return
+	}
+
+	// Verify task exists before reading the potentially large payload.
+	// This prevents DoS attacks where attackers spam large payloads for
+	// non-existent or deleted tasks, forcing us to read 64KB into memory
+	// and do expensive JSON operations before the database rejects it.
+	// The UpsertTaskSnapshot will re-fetch for RBAC validation, but this
+	// early check protects against malicious load.
+	task, err := api.Database.GetTaskByID(ctx, taskID)
+	if err != nil {
+		if httpapi.Is404Error(err) {
+			httpapi.ResourceNotFound(rw)
+			return
+		}
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error fetching task.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	// Reject deleted tasks early.
+	if task.DeletedAt.Valid {
+		httpapi.ResourceNotFound(rw)
+		return
+	}
+
+	// Verify task belongs to this agent's workspace.
+	if !task.WorkspaceID.Valid || task.WorkspaceID.UUID != latestBuild.WorkspaceID {
+		httpapi.ResourceNotFound(rw)
+		return
+	}
+
+	// Limit payload size to avoid excessive memory or data usage.
+	r.Body = http.MaxBytesReader(rw, r.Body, taskSnapshotMaxSize)
+
+	// Create envelope to store validated payload.
+	envelope := TaskLogSnapshotEnvelope{
+		Format: format,
+	}
+
+	switch format {
+	case "agentapi":
+		var payload agentapisdk.GetMessagesResponse
+		if err := json.NewDecoder(r.Body).Decode(&payload); err != nil {
+			httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+				Message: "Failed to decode request payload.",
+				Detail:  err.Error(),
+			})
+			return
+		}
+		// Verify messages field exists (can be empty array).
+		if payload.Messages == nil {
+			httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+				Message: "Invalid agentapi payload structure.",
+				Detail:  `Missing required "messages" field.`,
+			})
+			return
+		}
+		envelope.Data = payload
+	default:
+		// Defensive branch, we already validated "agentapi" format but may add
+		// more formats in the future.
+		httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Invalid format parameter.",
+			Detail:  fmt.Sprintf(`Only "agentapi" format is currently supported, got %q.`, format),
+		})
+		return
+	}
+
+	// Marshal envelope with validated payload in a single pass.
+	snapshotJSON, err := json.Marshal(envelope)
+	if err != nil {
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Failed to create snapshot envelope.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	// Upsert to database using agent's RBAC context.
+	err = api.Database.UpsertTaskSnapshot(ctx, database.UpsertTaskSnapshotParams{
+		TaskID:               task.ID,
+		LogSnapshot:          json.RawMessage(snapshotJSON),
+		LogSnapshotCreatedAt: dbtime.Time(api.Clock.Now()),
+	})
+	if err != nil {
+		if httpapi.IsUnauthorizedError(err) {
+			httpapi.ResourceNotFound(rw)
+			return
+		}
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error storing snapshot.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	api.Logger.Debug(ctx, "stored task log snapshot",
+		slog.F("task_id", task.ID),
+		slog.F("workspace_id", latestBuild.WorkspaceID),
+		slog.F("snapshot_size_bytes", len(snapshotJSON)))
+
+	rw.WriteHeader(http.StatusNoContent)
+}