feat(agent): add controlled task execution

2026-09-24 04:19:02 +08:00
parent d018bcae67
commit b31033e9ce
48 changed files with 14498 additions and 86 deletions
@@ -1,14 +1,20 @@
 import { type Router } from "express";
+import { randomUUID } from "node:crypto";
 import {
  UI_MESSAGE_STREAM_HEADERS,
  type UIMessage,
  type UIMessageChunk,
 } from "ai";
 import { z } from "zod";
+import { createExecutionPlanSnapshot, executionPlanNeedsTaskPublish, readExecutionPlans } from "./executionPlan.js";
+import { executionPlanSchema } from "./executionPlan.js";

 import { isSupportedModel } from "../chat/models.js";
 import { type ChatSessionBridge } from "../chat/sessionBridge.js";
-import { getLocalAgentContext } from "../context/localContext.js";
+import {
+  getLocalAgentContext,
+  getTrustedProductCredential,
+} from "../context/localContext.js";
 import { type FrontendActionCoordinator } from "../frontendAction/coordinator.js";
 import { type LearningOrchestrator } from "../learning/orchestrator.js";
 import { logger } from "../logger.js";
@@ -21,7 +27,14 @@ import {
 import {
  type SessionMetadataStore,
  type SessionRecord,
+  readPurpose,
 } from "../sessions/metadataStore.js";
+import {
+  bindOperationalTaskCredential,
+  clearOperationalTaskLease,
+  operationalTaskClient,
+  releaseOperationalTaskCredential,
+} from "../product/operational-task-client.js";
 import { type SessionTranscriptStore } from "../sessions/transcriptStore.js";
 import { type SessionUiStateStore } from "../sessions/uiStateStore.js";
 import type { UIEnvelopePayload } from "../uiEnvelope/types.js";
@@ -60,6 +73,10 @@ import {
  upsertBackendQuestion,
  upsertBackendTodoUpdate,
 } from "./chatUiState.js";
+import {
+  agentTriggerContextSchema,
+  buildAgentTriggerPrompt,
+} from "./agentTrigger.js";

 const payloadSchema = z.object({
  id: z.string().max(128).optional(),
@@ -71,6 +88,7 @@ const payloadSchema = z.object({
  }).optional(),
  approval_mode: z.enum(["request", "always"]).optional().default("request"),
  capabilities: z.array(z.string().max(64)).max(20).optional().default([]),
+  trigger_context: agentTriggerContextSchema.optional(),
 });

 const toSessionUiStateContext = (sessionRecord: SessionRecord) => ({
@@ -121,6 +139,15 @@ const readLatestUserText = (messages: UIMessage[]) => {
 const toUiMessages = (messages: unknown[] | undefined) =>
  Array.isArray(messages) ? messages.filter(isUiMessage) : [];

+export const toCollaborationMessages = (messages: UIMessage[]) =>
+  messages.flatMap((message) => {
+    if (message.role !== "user" && message.role !== "assistant") return [];
+    const textParts = message.parts.filter(
+      (part) => isObjectRecord(part) && part.type === "text" && typeof part.text === "string",
+    );
+    return textParts.length > 0 ? [{ ...message, parts: textParts } as UIMessage] : [];
+  });
+
 const toLegacyMessages = (messages: UIMessage[]) =>
  messages.flatMap((message) => {
    const content = readUiMessageText(message);
@@ -142,6 +169,66 @@ const createFallbackUserUiMessage = (content: string): UIMessage => ({
  parts: [{ type: "text", text: content }],
 });

+export const buildFailedTaskExecutionPlan = (history: Record<string, unknown>[]) => {
+  const previous = history.at(-1);
+  const parsed = previous ? executionPlanSchema.safeParse(previous) : null;
+  if (!parsed?.success) {
+    return {
+      title: "工况任务执行失败",
+      reason: "Agent 模型或工具执行失败，已记录失败并停止本次执行。",
+      nodes: [{
+        id: "AGENT_RUN_FAILURE",
+        title: "Agent 执行失败",
+        status: "failed" as const,
+        owner: "agent" as const,
+        detail: "模型或工具执行失败，等待人工检查或重新发起。",
+        evidence: [],
+        action_refs: [],
+      }],
+      edges: [],
+    };
+  }
+  let changed = false;
+  const nodes = parsed.data.nodes.map((node) => {
+    if (node.status === "running") {
+      changed = true;
+      return {
+        ...node,
+        status: "failed" as const,
+        detail: "模型或工具执行失败，当前节点未完成，等待人工检查或重新发起。",
+      };
+    }
+    if (node.status === "pending" || node.status === "waiting") {
+      return {
+        ...node,
+        status: "cancelled" as const,
+        detail: `${node.detail ? `${node.detail}；` : ""}本次 Agent 执行失败，后续工作未启动。`,
+      };
+    }
+    return node;
+  });
+  if (!changed && !nodes.some((node) => node.status === "failed")) {
+    const known = new Set(nodes.map((node) => node.id));
+    let index = 1;
+    let id = "AGENT_RUN_FAILURE";
+    while (known.has(id)) id = `AGENT_RUN_FAILURE_${index++}`;
+    nodes.push({
+      id,
+      title: "Agent 执行失败",
+      status: "failed",
+      owner: "agent",
+      detail: "模型或工具执行失败，等待人工检查或重新发起。",
+      evidence: [],
+      action_refs: [],
+    });
+  }
+  return {
+    ...parsed.data,
+    reason: "Agent 模型或工具执行失败，已记录失败并停止本次执行。",
+    nodes,
+  };
+};
+
 const pipeUiChunk = (
  res: {
    write: (chunk: string) => void;
@@ -182,6 +269,7 @@ export const registerChatStreamRoute = (
  }: RegisterChatStreamRouteOptions,
 ) => {
  const { activeRuns, lastRunStatuses } = runState;
+  const startingSessions = new Set<string>();
  chatRouter.post("/stream", async (req, res) => {
    const parsed = payloadSchema.safeParse(req.body);
    if (!parsed.success) {
@@ -192,6 +280,8 @@ export const registerChatStreamRoute = (
      return;
    }

+    let boundOperationalSessionId: string | undefined;
+    let startingSessionId: string | undefined;
    try {
      const authContext = getLocalAgentContext(req);
      const projectId = authContext.projectId;
@@ -199,15 +289,25 @@ export const registerChatStreamRoute = (
      const userId = authContext.userId;
      const actorKey = authContext.actorKey;
      const projectKey = authContext.projectKey;
+      const trigger = parsed.data.trigger_context;
      const requestUiMessages = toUiMessages(parsed.data.messages);
-      const promptText =
+      const requestPromptText =
        parsed.data.message?.trim() || readLatestUserText(requestUiMessages);
+      const promptText = trigger
+        ? buildAgentTriggerPrompt(trigger)
+        : requestPromptText;
      if (!promptText) {
        res.status(400).json({
          message: "message or messages with latest user text is required",
        });
        return;
      }
+      if (trigger && !parsed.data.session_id?.trim()) {
+        res.status(400).json({
+          message: "event trigger requires a reserved session_id",
+        });
+        return;
+      }
      const requestedSessionId = resolveRequestedStreamSessionId({
        sessionId: parsed.data.session_id,
        promptText,
@@ -223,6 +323,19 @@ export const registerChatStreamRoute = (
        res.status(404).json({ message: "session not found" });
        return;
      }
+      if (trigger && existingSessionRecord?.triggerId !== trigger.trigger_id) {
+        res.status(409).json({ message: "trigger does not match session" });
+        return;
+      }
+      if (
+        existingSessionRecord &&
+        readPurpose(existingSessionRecord) === "task_execution" &&
+        existingSessionRecord.modelId &&
+        parsed.data.model !== existingSessionRecord.modelId
+      ) {
+        res.status(409).json({ message: "task execution model does not match the reserved session" });
+        return;
+      }
      const hadExistingRuntimeSession = Boolean(existingSessionRecord);

      const { binding, requestContext, created } = await sessionBridge.resolve({
@@ -237,6 +350,8 @@ export const registerChatStreamRoute = (
      setRuntimeSessionContext({
        ...(getRuntimeSessionContext(binding.sessionId)!),
        capabilities: parsed.data.capabilities,
+        purpose: existingSessionRecord ? readPurpose(existingSessionRecord) : "ordinary_chat",
+        taskId: existingSessionRecord?.taskId,
      });
      const { record: ensuredSessionRecord, created: sessionCreated } =
        await sessionMetadataStore.ensure({
@@ -247,6 +362,31 @@ export const registerChatStreamRoute = (
          userId,
        });
      const activeSessionRecord = await sessionMetadataStore.touch(ensuredSessionRecord);
+      const sessionPurpose = readPurpose(activeSessionRecord);
+      if (
+        activeRuns.get(activeSessionRecord.sessionId)?.status === "running" ||
+        startingSessions.has(activeSessionRecord.sessionId)
+      ) {
+        res.status(409).json({
+          message: "session is already streaming",
+          session_id: activeSessionRecord.sessionId,
+        });
+        return;
+      }
+      startingSessions.add(activeSessionRecord.sessionId);
+      startingSessionId = activeSessionRecord.sessionId;
+      bindOperationalTaskCredential(
+        binding.sessionId,
+        getTrustedProductCredential(req),
+      );
+      boundOperationalSessionId = binding.sessionId;
+      const resolvedTaskContext = activeSessionRecord.taskId
+        && (sessionPurpose === "task_collaboration" || sessionPurpose === "task_execution")
+        ? await operationalTaskClient.getAgentContext(binding.sessionId, activeSessionRecord.taskId)
+        : null;
+      const taskContext = sessionPurpose === "task_collaboration"
+        ? resolvedTaskContext
+        : null;
      const hasRuntimeConversation = hadExistingRuntimeSession
        ? await runtimeHasConversation(runtime, binding.sessionId)
        : false;
@@ -264,22 +404,22 @@ export const registerChatStreamRoute = (
        toSessionUiStateContext(activeSessionRecord),
      );
      const persistedUiMessages = toUiMessages(initialSessionState?.messages);
+      let executionPlanHistory = initialSessionState?.executionPlans ?? readExecutionPlans(initialSessionState?.messages ?? []);
+      const visiblePersistedMessages = sessionPurpose === "task_collaboration"
+        ? toCollaborationMessages(persistedUiMessages)
+        : persistedUiMessages;
+      const visibleRequestMessages = sessionPurpose === "task_collaboration"
+        ? toCollaborationMessages(requestUiMessages)
+        : requestUiMessages;
      const requestMessages =
-        requestUiMessages.length > 0
-          ? requestUiMessages
+        visibleRequestMessages.length > 0
+          ? visibleRequestMessages
          : [
-              ...persistedUiMessages,
+              ...visiblePersistedMessages,
              createFallbackUserUiMessage(promptText),
            ];
      let latestUiMessages = requestMessages;
      const baseMessages = toLegacyMessages(requestMessages);
-      if (activeRuns.get(activeSessionRecord.sessionId)?.status === "running") {
-        res.status(409).json({
-          message: "session is already streaming",
-          session_id: activeSessionRecord.sessionId,
-        });
-        return;
-      }
      const recentTurns = await sessionTranscriptStore.getRecentTurns(historyContext, 8);

      logger.info(
@@ -313,13 +453,17 @@ export const registerChatStreamRoute = (
      const activeRun: ActiveRun = {
        clientSessionId,
        controller: abortController,
-        messages: initialMessages,
+        messages: updateLastAssistantMessage(initialMessages, (message) => ({
+          ...message, executionPlans: executionPlanHistory,
+        })),
        pendingPermissions: new Map(),
        pendingQuestions: new Map(),
        status: "running",
        subscribers: new Set(),
      };
      activeRuns.set(clientSessionId, activeRun);
+      startingSessions.delete(activeSessionRecord.sessionId);
+      startingSessionId = undefined;
      lastRunStatuses.set(clientSessionId, "running");
      const sessionUiStateContext = toSessionUiStateContext(activeSessionRecord);
      const assistantUiMessageId = `assistant-${Date.now().toString(36)}`;
@@ -426,15 +570,17 @@ export const registerChatStreamRoute = (
      };
      writeUiChunk({ type: "start", messageId: assistantUiMessageId });
      let persistQueue = sessionUiStateStore.write(sessionUiStateContext, {
+        executionPlans: executionPlanHistory,
        sessionId: activeSessionRecord.sessionId,
        isTitleManuallyEdited: initialSessionState?.isTitleManuallyEdited ?? false,
-        messages: latestUiMessages,
+        messages: sessionPurpose === "task_collaboration" ? toCollaborationMessages(latestUiMessages) : latestUiMessages,
      });
      const queueSessionUiStatePersist = () => {
        const snapshot = {
+          executionPlans: executionPlanHistory,
          sessionId: activeSessionRecord.sessionId,
          isTitleManuallyEdited: initialSessionState?.isTitleManuallyEdited ?? false,
-          messages: latestUiMessages,
+          messages: sessionPurpose === "task_collaboration" ? toCollaborationMessages(latestUiMessages) : latestUiMessages,
        };
        persistQueue = persistQueue
          .catch((error) => {
@@ -486,9 +632,10 @@ export const registerChatStreamRoute = (

      req.on("close", handleClientClose);
      res.on("close", handleClientClose);
+      let taskArtifactPublishQueue: Promise<unknown> = Promise.resolve();

      const publish = (event: string, data: Record<string, unknown>) => {
-        const subscriberData =
+        let subscriberData =
          event === "token"
            ? {
                ...data,
@@ -610,6 +757,36 @@ export const registerChatStreamRoute = (
              }),
            );
          }
+        } else if (event === "execution_plan") {
+          try {
+            const snapshot = createExecutionPlanSnapshot(data, [{ executionPlans: executionPlanHistory }], clientSessionId);
+            executionPlanHistory = [...executionPlanHistory, snapshot];
+            subscriberData = snapshot;
+            activeRun.messages = updateLastAssistantMessage(activeRun.messages, (message) => ({
+              ...message,
+              executionPlans: [...(Array.isArray(message.executionPlans) ? message.executionPlans : []), snapshot],
+            }));
+            if (activeSessionRecord.taskId && executionPlanNeedsTaskPublish(data)) {
+              taskArtifactPublishQueue = taskArtifactPublishQueue
+                .catch(() => undefined)
+                .then(() => operationalTaskClient.publishArtifacts(binding.sessionId, activeSessionRecord.taskId!, {
+                  requestId: randomUUID(),
+                  sourceSessionId: clientSessionId,
+                  executionPlan: snapshot,
+                }))
+                .catch((error) => {
+                  logger.warn({ err: error, sessionId: clientSessionId, taskId: activeSessionRecord.taskId }, "failed to publish execution plan to task record");
+                });
+            }
+          } catch (error) {
+            logger.warn({ err: error, sessionId: clientSessionId }, "rejected invalid execution plan");
+            publish("progress", {
+              id: "execution-plan-rejected", phase: "execution_plan", status: "error",
+              title: "计划更新未通过校验，保留上一版",
+              detail: error instanceof Error ? error.message : "计划结构无效",
+            });
+            return;
+          }
        } else if (event === "todo_update") {
          const payload = data as TodoUpdatePayload;
          activeRun.messages = updateLastAssistantMessage(activeRun.messages, (message) => ({
@@ -629,8 +806,17 @@ export const registerChatStreamRoute = (
          }));
        }

-        for (const subscriber of activeRun.subscribers) {
-          subscriber.write(event, subscriberData);
+        const hiddenCollaborationEvent = sessionPurpose === "task_collaboration" && [
+          "progress",
+          "todo_update",
+          "tool_call",
+          "ui_envelope",
+          "execution_plan",
+        ].includes(event);
+        if (!hiddenCollaborationEvent) {
+          for (const subscriber of activeRun.subscribers) {
+            subscriber.write(event, subscriberData);
+          }
        }
        void queueSessionUiStatePersist().catch((error) => {
          logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
@@ -640,19 +826,62 @@ export const registerChatStreamRoute = (
      const unregisterFrontendActionSink = parsed.data.capabilities.includes("frontend-action@1")
        ? frontendActionCoordinator.registerSink(clientSessionId, (request) => publish("frontend_action", request as unknown as Record<string, unknown>))
        : undefined;
+      let taskHeartbeat: ReturnType<typeof setInterval> | undefined;
+      let taskHeartbeatActive = true;
+      let taskHeartbeatPromise: Promise<void> | undefined;
+      let taskLeaseFailure: Error | undefined;
+      let taskRunFinalized = false;
+      const stopTaskHeartbeat = async () => {
+        taskHeartbeatActive = false;
+        if (taskHeartbeat) clearInterval(taskHeartbeat);
+        await taskHeartbeatPromise?.catch(() => undefined);
+      };
+      if (sessionPurpose === "task_execution" && activeSessionRecord.taskId) {
+        taskHeartbeat = setInterval(() => {
+          if (!taskHeartbeatActive || taskHeartbeatPromise) return;
+          const heartbeat = operationalTaskClient.heartbeatTask(
+            binding.sessionId,
+            activeSessionRecord.taskId!,
+            activeSessionRecord.runId!,
+            clientSessionId,
+          ).then(() => undefined).catch((error) => {
+            if (!taskHeartbeatActive || abortController.signal.aborted) return;
+            taskLeaseFailure = error instanceof Error
+              ? error
+              : new Error("task lease heartbeat failed");
+            taskHeartbeatActive = false;
+            if (taskHeartbeat) clearInterval(taskHeartbeat);
+            logger.warn(
+              { err: error, sessionId: clientSessionId, taskId: activeSessionRecord.taskId },
+              "task lease heartbeat failed; stopping stale execution",
+            );
+            abortController.abort();
+          }).finally(() => {
+            if (taskHeartbeatPromise === heartbeat) taskHeartbeatPromise = undefined;
+          });
+          taskHeartbeatPromise = heartbeat;
+        }, 20_000);
+      }
+      taskHeartbeat?.unref();

      try {
-        const preparedMessage = await buildPromptWithLearningContext(
-          memoryStore,
-          requestContext.actorKey,
-          requestContext.projectKey,
-          {
-            recentTurns,
-            persistedMessages: baseMessages,
-            message: promptText,
-            restoreConversation: shouldRestoreConversation,
-          },
-        );
+        const preparedMessage = taskContext
+          ? [
+              "你正在与用户协作理解一个工况任务。只根据下面由后端授权返回的任务、计划、报告和证据快照解释判断；不要展示后台执行会话、工具轨迹或内部推理。若资料不足，明确说明缺口。用户要求变更任务时使用结构化工具写回。",
+              `任务快照：${JSON.stringify(taskContext)}`,
+              `用户问题：${promptText}`,
+            ].join("\n\n")
+          : await buildPromptWithLearningContext(
+              memoryStore,
+              requestContext.actorKey,
+              requestContext.projectKey,
+              {
+                recentTurns,
+                persistedMessages: baseMessages,
+                message: promptText,
+                restoreConversation: shouldRestoreConversation,
+              },
+            );
        const streamResult = await streamPromptResponse({
          runtime,
          sessionId: binding.sessionId,
@@ -668,6 +897,63 @@ export const registerChatStreamRoute = (
            publish(event, data);
          },
        });
+        if (taskLeaseFailure) {
+          publish("error", {
+            session_id: clientSessionId,
+            message: "工况任务租约已失效，本次执行已停止，等待任务重新领取。",
+            detail: taskLeaseFailure.message,
+          });
+        } else if (streamResult.failed && sessionPurpose === "task_execution" && activeSessionRecord.taskId) {
+          publish("execution_plan", buildFailedTaskExecutionPlan(executionPlanHistory));
+        }
+        await taskArtifactPublishQueue;
+        await stopTaskHeartbeat();
+        if (
+          sessionPurpose === "task_execution"
+          && activeSessionRecord.taskId
+          && activeSessionRecord.runId
+          && !taskLeaseFailure
+        ) {
+          if (streamResult.failed || streamResult.aborted) {
+            await operationalTaskClient.failTaskRun(
+              binding.sessionId,
+              activeSessionRecord.taskId,
+              activeSessionRecord.runId,
+              clientSessionId,
+              streamResult.aborted ? "AGENT_ABORTED" : "AGENT_STREAM_FAILED",
+              streamResult.aborted ? "Agent execution was aborted" : "Agent stream failed",
+            );
+            taskRunFinalized = true;
+          } else {
+            try {
+              const completedTask = await operationalTaskClient.completeTaskRun(
+                binding.sessionId,
+                activeSessionRecord.taskId,
+                activeSessionRecord.runId,
+                clientSessionId,
+              );
+              taskRunFinalized = true;
+              if (completedTask.lifecycle_status === "failed") {
+                throw new Error("Server rejected the run result because required artifacts were incomplete");
+              }
+            } catch (error) {
+              if (!taskRunFinalized) {
+                const code = error && typeof error === "object" && "code" in error
+                  && typeof error.code === "string" ? error.code : "AGENT_COMPLETION_FAILED";
+                await operationalTaskClient.failTaskRun(
+                  binding.sessionId,
+                  activeSessionRecord.taskId,
+                  activeSessionRecord.runId,
+                  clientSessionId,
+                  code,
+                  error instanceof Error ? error.message : "Agent run completion failed",
+                );
+                taskRunFinalized = true;
+              }
+              throw error;
+            }
+          }
+        }
        await persistQueue.catch((error) => {
          logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
        });
@@ -720,7 +1006,7 @@ export const registerChatStreamRoute = (
            });
          }
          const latestTurn = extractLatestFrontendTurn(activeRun.messages);
-          if (latestTurn) {
+          if (latestTurn && sessionPurpose === "ordinary_chat") {
            void learningOrchestrator.onTurnCompleted({
              ...latestTurn,
              requestContext,
@@ -733,9 +1019,37 @@ export const registerChatStreamRoute = (
            });
          }
        }
+      } catch (error) {
+        await stopTaskHeartbeat();
+        if (
+          sessionPurpose === "task_execution" &&
+          activeSessionRecord.taskId
+        ) {
+          publish("execution_plan", buildFailedTaskExecutionPlan(executionPlanHistory));
+          await taskArtifactPublishQueue;
+          if (activeSessionRecord.runId && !taskLeaseFailure && !taskRunFinalized) {
+            await operationalTaskClient.failTaskRun(
+              binding.sessionId,
+              activeSessionRecord.taskId,
+              activeSessionRecord.runId,
+              clientSessionId,
+              abortController.signal.aborted ? "AGENT_ABORTED" : "AGENT_EXECUTION_ERROR",
+              abortController.signal.aborted
+                ? "Agent execution was aborted"
+                : error instanceof Error ? error.message : "Agent execution failed",
+            ).then(() => {
+              taskRunFinalized = true;
+            }).catch((finishError) => {
+              logger.warn({ err: finishError, sessionId: clientSessionId }, "failed to record task run failure");
+            });
+          }
+        }
+        throw error;
      } finally {
+        await stopTaskHeartbeat();
+        await taskArtifactPublishQueue.catch(() => undefined);
        unregisterFrontendActionSink?.();
-        if (abortController.signal.aborted) {
+        if (abortController.signal.aborted && !taskLeaseFailure) {
          activeRun.messages = updateLastAssistantMessage(activeRun.messages, (message) => ({
            ...message,
            content:
@@ -764,10 +1078,11 @@ export const registerChatStreamRoute = (
          logger.warn({ err: error, sessionId: clientSessionId }, "failed to persist chat stream state");
        });
        sessionBridge.finalizeRequest(clientSessionId);
-        activeRun.status = abortController.signal.aborted
-          ? activeRun.status === "aborted"
-            ? "aborted"
-            : "aborted"
+        releaseOperationalTaskCredential(binding.sessionId);
+        activeRun.status = taskLeaseFailure
+          ? "error"
+          : abortController.signal.aborted
+          ? "aborted"
          : activeRun.status === "running"
            ? "completed"
            : activeRun.status;
@@ -791,6 +1106,12 @@ export const registerChatStreamRoute = (
        res.end();
      }
    } catch (error) {
+      if (startingSessionId) {
+        startingSessions.delete(startingSessionId);
+      }
+      if (boundOperationalSessionId) {
+        releaseOperationalTaskCredential(boundOperationalSessionId);
+      }
      const detail = error instanceof Error ? error.message : String(error);
      logger.error({ err: error }, "chat stream failed");
      if (res.headersSent) {