diff --git a/apps/desktop-client/src/main/runtime-controller.ts b/apps/desktop-client/src/main/runtime-controller.ts index a50c14c..6182a26 100644 --- a/apps/desktop-client/src/main/runtime-controller.ts +++ b/apps/desktop-client/src/main/runtime-controller.ts @@ -749,6 +749,7 @@ async function sendHeartbeat(source: "startup" | "timer"): Promise { cpu_arch: process.arch, client_version: state.client?.client_version ?? "0.1.0-dev", channel: state.client?.channel ?? "dev", + startup: source === "startup", account_ids: state.accounts .filter((account) => getProjectedAccountHealth({ @@ -1634,18 +1635,17 @@ async function executeLeasedDesktopMonitorTask( buildMonitoringResultPayload(leased.lease_token as string, execution), ); - const completed = await completeDesktopTask(taskRecord.id, { - lease_token: leased.lease_token as string, - status: execution.status, - payload: execution.payload, - error: execution.error, - }); - upsertTaskFromInfo(completed, { - routing, - summary: execution.summary, - attemptId: null, - leaseToken: null, - }); + const existing = state.tasks.get(taskRecord.id); + if (existing) { + existing.status = execution.status; + existing.summary = execution.summary; + existing.result = execution.payload ?? null; + existing.error = execution.error ?? null; + existing.updatedAt = Date.now(); + existing.attemptId = null; + existing.leaseToken = null; + state.tasks.set(taskRecord.id, existing); + } noteLeaseReleased(execution.status === "failed" ? "failed" : "completed", taskRecord.id); noteMonitorTaskCompleted(taskRecord.id); recordActivity( @@ -1738,29 +1738,15 @@ async function executeLeasedDesktopMonitorTask( const failureSummary = monitoringResultWriteError ? `${failureMessage || "监控任务执行失败。"}（monitoring 失败回传失败：${errorMessage(monitoringResultWriteError)}）` : (failureMessage || "监控任务执行失败。"); - const completed = await completeDesktopTask(taskRecord.id, { - lease_token: leased.lease_token as string, - status: "failed", - error: structuredError, - }).catch(async (completeError) => { - const existing = state.tasks.get(taskRecord.id); - if (existing) { - existing.status = "failed"; - existing.error = structuredError; - existing.summary = `${failureSummary}（desktop task 回写失败：${errorMessage(completeError)}）`; - existing.updatedAt = Date.now(); - state.tasks.set(taskRecord.id, existing); - } - return null; - }); - - if (completed) { - upsertTaskFromInfo(completed, { - routing, - summary: failureSummary, - attemptId: null, - leaseToken: null, - }); + const existing = state.tasks.get(taskRecord.id); + if (existing) { + existing.status = "failed"; + existing.error = structuredError; + existing.summary = failureSummary; + existing.updatedAt = Date.now(); + existing.attemptId = null; + existing.leaseToken = null; + state.tasks.set(taskRecord.id, existing); } noteLeaseReleased("failed", taskRecord.id); noteMonitorTaskCompleted(taskRecord.id); diff --git a/packages/shared-types/src/index.ts b/packages/shared-types/src/index.ts index f780692..7c87579 100644 --- a/packages/shared-types/src/index.ts +++ b/packages/shared-types/src/index.ts @@ -104,6 +104,7 @@ export interface DesktopClientHeartbeatRequest { client_version?: string; channel?: string; account_ids?: string[]; + startup?: boolean; } export interface DesktopClientHeartbeatResponse { diff --git a/server/internal/tenant/app/desktop_client_service.go b/server/internal/tenant/app/desktop_client_service.go index b97c0fe..0ca7285 100644 --- a/server/internal/tenant/app/desktop_client_service.go +++ b/server/internal/tenant/app/desktop_client_service.go @@ -21,14 +21,20 @@ import ( const DesktopClientTokenTTL = 30 * 24 * time.Hour type DesktopClientService struct { - repo repository.DesktopClientRepository - redis *goredis.Client + repo repository.DesktopClientRepository + redis *goredis.Client + taskSvc *DesktopTaskService } func NewDesktopClientService(repo repository.DesktopClientRepository, redis *goredis.Client) *DesktopClientService { return &DesktopClientService{repo: repo, redis: redis} } +func (s *DesktopClientService) WithTaskService(taskSvc *DesktopTaskService) *DesktopClientService { + s.taskSvc = taskSvc + return s +} + type RegisterDesktopClientRequest struct { ClientID string `json:"client_id"` DeviceName string `json:"device_name" binding:"required"` @@ -45,6 +51,7 @@ type HeartbeatDesktopClientRequest struct { ClientVersion *string `json:"client_version"` Channel *string `json:"channel"` AccountIDs []string `json:"account_ids"` + Startup bool `json:"startup"` } type DesktopClientView struct { @@ -182,6 +189,12 @@ func (s *DesktopClientService) Heartbeat(ctx context.Context, client *repository markDesktopClientPresent(ctx, s.redis, updated.ID) markDesktopAccountsPresent(ctx, s.redis, updated.ID, parseDesktopAccountIDs(req.AccountIDs)) + if req.Startup && s.taskSvc != nil { + if err := s.taskSvc.recoverClientTasksOnStartup(ctx, updated); err != nil { + return nil, err + } + } + return &HeartbeatDesktopClientResponse{ Client: buildDesktopClientView(updated), ServerTime: time.Now().UTC(), @@ -202,6 +215,11 @@ func (s *DesktopClientService) Revoke(ctx context.Context, client *repository.De } clearDesktopClientPresenceState(ctx, s.redis, updated.ID) + if s.taskSvc != nil { + if err := s.taskSvc.recoverClientTasksOnDisconnect(ctx, updated); err != nil { + return nil, err + } + } view := buildDesktopClientView(updated) return &view, nil @@ -213,6 +231,11 @@ func (s *DesktopClientService) Offline(ctx context.Context, client *repository.D } clearDesktopClientPresenceState(ctx, s.redis, client.ID) + if s.taskSvc != nil { + if err := s.taskSvc.recoverClientTasksOnDisconnect(ctx, client); err != nil { + return nil, err + } + } return &OfflineDesktopClientResponse{ ServerTime: time.Now().UTC(), diff --git a/server/internal/tenant/app/desktop_task_service.go b/server/internal/tenant/app/desktop_task_service.go index 859cb03..3cb0e9e 100644 --- a/server/internal/tenant/app/desktop_task_service.go +++ b/server/internal/tenant/app/desktop_task_service.go @@ -132,6 +132,9 @@ func (s *DesktopTaskService) Lease(ctx context.Context, client *repository.Deskt if err := s.dropStaleMonitorTasksForClient(ctx, client); err != nil { return nil, err } + if err := s.recoverExpiredClientTasks(ctx, client); err != nil { + return nil, err + } rawToken, tokenHash, err := newDesktopClientToken() if err != nil { @@ -783,6 +786,245 @@ func normalizeDesktopTaskTerminalStatus(status string) string { } } +type desktopTaskRecoveryMode string + +const ( + desktopTaskRecoveryModeStartup desktopTaskRecoveryMode = "startup" + desktopTaskRecoveryModeDisconnect desktopTaskRecoveryMode = "disconnect" + desktopTaskRecoveryModeLeaseExpiry desktopTaskRecoveryMode = "lease_expired" +) + +type recoveredDesktopTaskLease struct { + TaskID uuid.UUID + WorkspaceID int64 + Kind string + Status string + ActiveAttempt pgtype.UUID +} + +func (s *DesktopTaskService) recoverClientTasksOnStartup(ctx context.Context, client *repository.DesktopClient) error { + return s.recoverClientTasks(ctx, client, desktopTaskRecoveryModeStartup) +} + +func (s *DesktopTaskService) recoverClientTasksOnDisconnect(ctx context.Context, client *repository.DesktopClient) error { + return s.recoverClientTasks(ctx, client, desktopTaskRecoveryModeDisconnect) +} + +func (s *DesktopTaskService) recoverExpiredClientTasks(ctx context.Context, client *repository.DesktopClient) error { + return s.recoverClientTasks(ctx, client, desktopTaskRecoveryModeLeaseExpiry) +} + +func (s *DesktopTaskService) recoverClientTasks( + ctx context.Context, + client *repository.DesktopClient, + mode desktopTaskRecoveryMode, +) error { + if s == nil || client == nil || s.pool == nil { + return nil + } + + publishErrorJSON, publishReason, publishMessage, err := desktopTaskRecoveryPayload(mode, "publish") + if err != nil { + return response.ErrInternal(50195, "desktop_task_recovery_payload_failed", "failed to encode publish desktop task recovery payload") + } + monitorErrorJSON, monitorReason, _, err := desktopTaskRecoveryPayload(mode, "monitor") + if err != nil { + return response.ErrInternal(50195, "desktop_task_recovery_payload_failed", "failed to encode monitor desktop task recovery payload") + } + + tx, err := s.pool.Begin(ctx) + if err != nil { + return response.ErrInternal(50196, "desktop_task_recovery_begin_failed", "failed to start desktop task recovery") + } + defer tx.Rollback(ctx) + + repo := repository.NewDesktopTaskRepository(tx) + query := ` + SELECT + desktop_id, + workspace_id, + kind, + active_attempt_id + FROM desktop_tasks + WHERE target_client_id = $1 + AND workspace_id = $2 + AND status = 'in_progress' + ` + queryArgs := []any{client.ID, client.WorkspaceID} + if mode == desktopTaskRecoveryModeLeaseExpiry { + query += ` + AND lease_expires_at IS NOT NULL + AND lease_expires_at < NOW() + ` + } + query += ` + FOR UPDATE + ` + + rows, err := tx.Query(ctx, query, queryArgs...) + if err != nil { + s.logWarn("desktop task recovery query failed", err, zap.String("client_id", client.ID.String()), zap.String("mode", string(mode))) + return response.ErrInternal(50197, "desktop_task_recovery_query_failed", "failed to select recoverable desktop tasks") + } + defer rows.Close() + + recovered := make([]recoveredDesktopTaskLease, 0) + for rows.Next() { + var item recoveredDesktopTaskLease + if scanErr := rows.Scan( + &item.TaskID, + &item.WorkspaceID, + &item.Kind, + &item.Status, + &item.ActiveAttempt, + ); scanErr != nil { + return response.ErrInternal(50198, "desktop_task_recovery_scan_failed", "failed to parse recovered desktop tasks") + } + recovered = append(recovered, item) + } + if err := rows.Err(); err != nil { + return response.ErrInternal(50198, "desktop_task_recovery_scan_failed", "failed to iterate recovered desktop tasks") + } + + for index := range recovered { + item := &recovered[index] + if item.Kind == "monitor" { + if _, execErr := tx.Exec(ctx, ` + UPDATE desktop_tasks + SET status = 'queued', + error = $2, + active_attempt_id = NULL, + lease_token_hash = NULL, + lease_expires_at = NULL, + interrupted_at = COALESCE(interrupted_at, NOW()), + interrupt_reason = $3, + enqueued_at = NOW(), + updated_at = NOW() + WHERE desktop_id = $1 + `, item.TaskID, monitorErrorJSON, monitorReason); execErr != nil { + s.logWarn("desktop monitor task recovery update failed", execErr, zap.String("task_id", item.TaskID.String()), zap.String("mode", string(mode))) + return response.ErrInternal(50197, "desktop_task_recovery_query_failed", "failed to select recoverable desktop tasks") + } + item.Status = "queued" + continue + } + + if _, execErr := tx.Exec(ctx, ` + UPDATE desktop_tasks + SET status = 'unknown', + error = $2, + active_attempt_id = NULL, + lease_token_hash = NULL, + lease_expires_at = NULL, + interrupted_at = COALESCE(interrupted_at, NOW()), + interrupt_reason = COALESCE(interrupt_reason, $3), + updated_at = NOW() + WHERE desktop_id = $1 + `, item.TaskID, publishErrorJSON, publishReason); execErr != nil { + s.logWarn("desktop publish task recovery update failed", execErr, zap.String("task_id", item.TaskID.String()), zap.String("mode", string(mode))) + return response.ErrInternal(50197, "desktop_task_recovery_query_failed", "failed to select recoverable desktop tasks") + } + item.Status = "unknown" + } + + for _, item := range recovered { + if !item.ActiveAttempt.Valid { + continue + } + + attemptID, convErr := uuid.FromBytes(item.ActiveAttempt.Bytes[:]) + if convErr != nil { + s.logWarn("desktop recovery attempt id decode failed", convErr, zap.String("task_id", item.TaskID.String())) + continue + } + + finalStatus := literalStringPtr(item.Status) + errorJSON := publishErrorJSON + if item.Kind == "monitor" { + finalStatus = literalStringPtr("aborted") + errorJSON = monitorErrorJSON + } + + if finishErr := repo.FinishAttempt(ctx, repository.FinishDesktopTaskAttemptParams{ + TaskID: item.TaskID, + AttemptID: attemptID, + FinalStatus: finalStatus, + Error: errorJSON, + }); finishErr != nil { + s.logWarn("desktop recovery attempt finish failed", finishErr, zap.String("task_id", item.TaskID.String())) + } + } + + if err := tx.Commit(ctx); err != nil { + return response.ErrInternal(50199, "desktop_task_recovery_commit_failed", "failed to commit desktop task recovery") + } + + for _, item := range recovered { + task, getErr := s.repo.GetByDesktopID(ctx, item.TaskID, item.WorkspaceID) + if getErr != nil { + s.logWarn("desktop recovery reload failed", getErr, zap.String("task_id", item.TaskID.String())) + continue + } + + eventType := "task_completed" + if item.Kind == "monitor" { + eventType = "task_available" + } + s.publishTaskEvent(ctx, task, eventType) + } + + if len(recovered) > 0 && s.logger != nil { + s.logger.Info("desktop tasks recovered", + zap.String("client_id", client.ID.String()), + zap.String("mode", string(mode)), + zap.Int("count", len(recovered)), + zap.String("publish_message", publishMessage), + ) + } + + return nil +} + +func desktopTaskRecoveryPayload(mode desktopTaskRecoveryMode, kind string) ([]byte, string, string, error) { + reason := strings.TrimSpace(string(mode)) + message := "desktop task was recovered after the client lost the active lease" + source := "desktop_task_recovery" + + switch mode { + case desktopTaskRecoveryModeStartup: + source = "desktop_client_startup" + if kind == "monitor" { + message = "desktop client restarted while the monitor task was in progress; task has been re-queued" + } else { + message = "desktop client restarted while the publish task was in progress; task has been moved to unknown for manual reconcile" + } + case desktopTaskRecoveryModeDisconnect: + source = "desktop_client_offline" + if kind == "monitor" { + message = "desktop client went offline while the monitor task was in progress; task has been re-queued" + } else { + message = "desktop client went offline while the publish task was in progress; task has been moved to unknown for manual reconcile" + } + case desktopTaskRecoveryModeLeaseExpiry: + source = "desktop_task_lease_expiry" + if kind == "monitor" { + message = "desktop task lease expired before monitor completion; task has been re-queued" + } else { + message = "desktop task lease expired before publish completion; task has been moved to unknown for manual reconcile" + } + } + + payload, err := marshalOptionalJSON(map[string]any{ + "reason": reason, + "message": message, + "source": source, + }) + if err != nil { + return nil, "", "", err + } + return payload, reason, message, nil +} + func (s *DesktopTaskService) dropStaleMonitorTasksForClient(ctx context.Context, client *repository.DesktopClient) error { if s == nil || client == nil || s.pool == nil || s.monitoringPool == nil { return nil diff --git a/server/internal/tenant/app/monitoring_callback_service.go b/server/internal/tenant/app/monitoring_callback_service.go index 81caeca..ebaae75 100644 --- a/server/internal/tenant/app/monitoring_callback_service.go +++ b/server/internal/tenant/app/monitoring_callback_service.go @@ -13,6 +13,7 @@ import ( "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgtype" "github.com/jackc/pgx/v5/pgxpool" "go.uber.org/zap" "golang.org/x/net/publicsuffix" @@ -415,6 +416,16 @@ func (s *MonitoringCallbackService) handleTaskResultForInstallation( QueueEnqueuedAt: &enqueuedAt, }) + if task.ExecutionOwner == "desktop_tasks" { + if finalizeErr := s.completeLinkedDesktopMonitorTask(ctx, task, req); finalizeErr != nil && s.logger != nil { + s.logger.Warn("phase2 monitor desktop task finalize failed", + zap.Int64("monitor_task_id", task.ID), + zap.String("ai_platform_id", task.AIPlatformID), + zap.Error(finalizeErr), + ) + } + } + return &MonitoringTaskResultResponse{ TaskID: task.ID, TaskStatus: "received", @@ -1757,6 +1768,190 @@ func (s *MonitoringCallbackService) processTaskResult(ctx context.Context, task }, nil } +func buildDesktopMonitorTaskResultValue(task *monitoringCollectTask, req MonitoringTaskResultRequest, runStatus string) map[string]any { + result := map[string]any{ + "monitoring_task_id": task.ID, + "status": runStatus, + } + + if task != nil { + if strings.TrimSpace(task.AIPlatformID) != "" { + result["platform"] = strings.TrimSpace(task.AIPlatformID) + } + if strings.TrimSpace(task.QuestionText) != "" { + result["question_text"] = strings.TrimSpace(task.QuestionText) + } + } + if value := strings.TrimSpace(optionalStringValue(req.ProviderModel)); value != "" { + result["provider_model"] = value + } + if value := strings.TrimSpace(optionalStringValue(req.ProviderRequestID)); value != "" { + result["provider_request_id"] = value + } + if value := strings.TrimSpace(optionalStringValue(req.RequestID)); value != "" { + result["request_id"] = value + } + if value := strings.TrimSpace(optionalStringValue(req.Answer)); value != "" { + result["answer"] = value + } + if len(req.Citations) > 0 { + result["citations"] = req.Citations + } + if len(req.SearchResults) > 0 { + result["search_results"] = req.SearchResults + } + if len(req.RawResponseJSON) > 0 { + result["raw_response_json"] = req.RawResponseJSON + } + return result +} + +func buildDesktopMonitorTaskErrorValue(task *monitoringCollectTask, req MonitoringTaskResultRequest) map[string]any { + message := strings.TrimSpace(optionalStringValue(req.ErrorMessage)) + if message == "" { + message = "monitoring task failed" + } + + result := map[string]any{ + "code": "monitoring_task_failed", + "message": message, + "monitoring_task_id": task.ID, + } + if task != nil && strings.TrimSpace(task.AIPlatformID) != "" { + result["platform"] = strings.TrimSpace(task.AIPlatformID) + } + if len(req.RawResponseJSON) > 0 { + result["raw_response_json"] = req.RawResponseJSON + } + return result +} + +func (s *MonitoringCallbackService) completeLinkedDesktopMonitorTask( + ctx context.Context, + task *monitoringCollectTask, + req MonitoringTaskResultRequest, +) error { + if s == nil || s.businessPool == nil || task == nil || task.ExecutionOwner != "desktop_tasks" { + return nil + } + + runStatus := normalizeRunStatus(req.Status) + if runStatus == "" { + return nil + } + + tx, err := s.businessPool.Begin(ctx) + if err != nil { + return response.ErrInternal(50041, "desktop_task_complete_begin_failed", "failed to start phase2 desktop task completion") + } + defer tx.Rollback(ctx) + + var ( + desktopID uuid.UUID + workspaceID int64 + activeAttemptID pgtype.UUID + ) + if err := tx.QueryRow(ctx, ` + SELECT desktop_id, workspace_id, active_attempt_id + FROM desktop_tasks + WHERE kind = 'monitor' + AND monitor_task_id = $1 + AND status IN ('in_progress', 'unknown') + ORDER BY created_at DESC, id DESC + LIMIT 1 + FOR UPDATE + `, task.ID).Scan(&desktopID, &workspaceID, &activeAttemptID); err != nil { + if err == pgx.ErrNoRows { + return nil + } + return response.ErrInternal(50041, "desktop_task_lookup_failed", "failed to inspect phase2 desktop task state") + } + + resultJSON, err := marshalOptionalJSON(buildDesktopMonitorTaskResultValue(task, req, runStatus)) + if err != nil { + return response.ErrInternal(50041, "desktop_task_payload_encode_failed", "failed to encode phase2 desktop task result payload") + } + + var errorJSON []byte + if runStatus == "failed" { + errorJSON, err = marshalOptionalJSON(buildDesktopMonitorTaskErrorValue(task, req)) + if err != nil { + return response.ErrInternal(50041, "desktop_task_payload_encode_failed", "failed to encode phase2 desktop task error payload") + } + } + + if _, err := tx.Exec(ctx, ` + UPDATE desktop_tasks + SET status = $2, + result = $3, + error = $4, + active_attempt_id = NULL, + lease_token_hash = NULL, + lease_expires_at = NULL, + updated_at = NOW() + WHERE desktop_id = $1 + `, desktopID, runStatus, resultJSON, errorJSON); err != nil { + return response.ErrInternal(50041, "desktop_task_complete_failed", "failed to finalize phase2 desktop task") + } + + taskRepo := repository.NewDesktopTaskRepository(tx) + if activeAttemptID.Valid { + resolvedAttemptID, convErr := uuid.FromBytes(activeAttemptID.Bytes[:]) + if convErr != nil { + return response.ErrInternal(50041, "desktop_task_attempt_decode_failed", "failed to decode phase2 desktop task attempt id") + } + if finishErr := taskRepo.FinishAttempt(ctx, repository.FinishDesktopTaskAttemptParams{ + TaskID: desktopID, + AttemptID: resolvedAttemptID, + FinalStatus: &runStatus, + Error: errorJSON, + }); finishErr != nil { + return response.ErrInternal(50041, "desktop_task_attempt_finish_failed", "failed to finalize phase2 desktop task attempt") + } + } + + finalizedTask, err := taskRepo.GetByDesktopID(ctx, desktopID, workspaceID) + if err != nil { + return response.ErrInternal(50041, "desktop_task_reload_failed", "failed to reload finalized phase2 desktop task") + } + + if err := tx.Commit(ctx); err != nil { + return response.ErrInternal(50041, "desktop_task_complete_commit_failed", "failed to commit phase2 desktop task completion") + } + + title, businessDate, schedulerGroupKey, questionText := deriveDesktopTaskEventMetadataFromPayload( + finalizedTask.Kind, + finalizedTask.Payload, + ) + if publishErr := publishDesktopTaskEvent(ctx, s.rabbitMQ, DesktopTaskEvent{ + Type: "task_completed", + TaskID: finalizedTask.DesktopID.String(), + JobID: finalizedTask.JobID.String(), + WorkspaceID: finalizedTask.WorkspaceID, + TargetAccountID: finalizedTask.TargetAccountID.String(), + TargetClientID: finalizedTask.TargetClientID.String(), + Platform: finalizedTask.Platform, + Title: title, + BusinessDate: businessDate, + SchedulerGroupKey: schedulerGroupKey, + QuestionText: questionText, + Status: finalizedTask.Status, + Kind: finalizedTask.Kind, + Priority: finalizedTask.Priority, + Lane: finalizedTask.Lane, + InterruptGeneration: finalizedTask.InterruptGeneration, + UpdatedAt: finalizedTask.UpdatedAt, + }); publishErr != nil && s.logger != nil { + s.logger.Warn("phase2 desktop task completion event publish failed", + zap.Int64("monitor_task_id", task.ID), + zap.String("desktop_task_id", finalizedTask.DesktopID.String()), + zap.Error(publishErr), + ) + } + + return nil +} + func (s *MonitoringCallbackService) replaceCitationFacts(ctx context.Context, tx pgx.Tx, tenantID, runID int64) error { if _, err := tx.Exec(ctx, ` DELETE FROM monitoring_citation_facts diff --git a/server/internal/tenant/transport/desktop_client_handler.go b/server/internal/tenant/transport/desktop_client_handler.go index d8a4aac..458382e 100644 --- a/server/internal/tenant/transport/desktop_client_handler.go +++ b/server/internal/tenant/transport/desktop_client_handler.go @@ -15,11 +15,17 @@ type DesktopClientHandler struct { } func NewDesktopClientHandler(a *bootstrap.App) *DesktopClientHandler { + taskSvc := app.NewDesktopTaskService( + a.DB, + a.MonitoringDB, + a.RabbitMQ, + a.Logger, + ).WithCache(a.Cache) return &DesktopClientHandler{ svc: app.NewDesktopClientService( repository.NewDesktopClientRepository(a.DB), a.Redis, - ), + ).WithTaskService(taskSvc), } }