Mujoco_WASM/web_platform/src/mobile/agent/AgentTaskPanel.tsx

import { useEffect, useRef, useState } from 'react';
import { Button } from '../../components/ui';
import { WEBSITE_MODE } from '../../app/websiteConfig';
import { downloadBytes } from '../../project/cachedFiles';
import type { SimulationSession } from '../../simulation/SimulationSession';
import { AgentTaskController } from './AgentTaskController';
import { DecisionClient } from './DecisionClient';
import { DeterministicBaseline } from './DeterministicBaseline';
import { MockDecisionProvider } from './MockDecisionProvider';
import { ModelConnectionPanel } from './ModelConnectionPanel';
import {
  PICK_PLACE,
  TASK_GOAL_LIMITS,
  seedPositions,
  type PickPlaceTaskConfig,
} from './PickPlaceScene';

export interface AgentWorkbenchBridge {
  loadScene(config: PickPlaceTaskConfig): Promise<SimulationSession>;
  getSession(): SimulationSession | undefined;
  isTaskScene(): boolean;
  publish(): void;
}
const states = {
  idle: '准备',
  planning: '等待 LLM 规划',
  deciding: '等待 Jev 判定',
  executing: '执行技能',
  paused: '任务已暂停',
  succeeded: '物理验收成功',
  failed: '失败 / 安全停止',
  cancelled: '已取消',
};
export function AgentTaskPanel({
  bridge,
  loading,
}: {
  bridge: AgentWorkbenchBridge;
  loading: boolean;
}) {
  const liveBridge = useRef(bridge);
  useEffect(() => {
    liveBridge.current = bridge;
  }, [bridge]);
  const task = useRef<AgentTaskController | undefined>(undefined);
  const ownSession = useRef<SimulationSession | undefined>(undefined);
  const [snapshot, setSnapshot] = useState<ReturnType<AgentTaskController['snapshot']>>();
  const [client, setClient] = useState<DecisionClient>();
  const [settingsOpen, setSettingsOpen] = useState(false);
  const [mode, setMode] = useState<'mock' | 'api'>('mock');
  const [instruction, setInstruction] = useState(
    '把方块抓起，使用底盘搬运到指定支撑台，放下、松开并撤离。',
  );
  const [goal, setGoal] = useState<number[]>([...PICK_PLACE.goal.position]);
  const [busy, setBusy] = useState(false);
  const [message, setMessage] = useState('需要 A / lekiwi-v1 资产；不会自动下载或替换机器人。');
  const epoch = useRef(0);
  const sceneDirty = useRef(true);
  const alive = useRef(true);
  function publish() {
    const current = task.current;
    if (current) setSnapshot(current.snapshot());
    liveBridge.current.publish();
  }
  function stop(reason = '用户停止任务') {
    epoch.current++;
    const current = task.current,
      session = ownSession.current;
    if (current) current.cancel(reason);
    if (
      session &&
      session === liveBridge.current.getSession() &&
      session.snapshot().controlOwner === 'agent'
    ) {
      session.stopAgent(reason);
      session.setPaused(true);
    }
    if (alive.current) publish();
  }
  useEffect(() => {
    alive.current = true;
    const timer = setInterval(() => {
      const current = task.current,
        session = ownSession.current;
      if (!current) return;
      setSnapshot(current.snapshot());
      if (
        current.terminal &&
        session === liveBridge.current.getSession() &&
        session?.snapshot().controlOwner === 'agent'
      ) {
        session.stopAgent('任务已结束');
        session.setPaused(true);
        liveBridge.current.publish();
      }
    }, 250);
    const invalidatePending = () => {
      epoch.current++;
    };
    return () => {
      alive.current = false;
      invalidatePending();
      clearInterval(timer);
      const session = ownSession.current;
      if (
        session === liveBridge.current.getSession() &&
        session?.snapshot().controlOwner === 'agent'
      ) {
        session.stopAgent('任务面板已卸载');
        session.setPaused(true);
      }
      task.current?.dispose();
    };
  }, []);
  async function load() {
    stop('任务场景重新加载');
    const generation = epoch.current;
    setBusy(true);
    try {
      seedPositions(0, goal);
      const session = await liveBridge.current.loadScene({ seed: 0, goal: [...goal] });
      if (generation !== epoch.current || !alive.current) return;
      ownSession.current = session;
      task.current = undefined;
      sceneDirty.current = false;
      setSnapshot(undefined);
      setMessage('A 抓放场景已加载并暂停；用户地图草稿保留，但不用于本次空旷平地任务。');
    } catch (error) {
      if (generation === epoch.current && alive.current)
        setMessage(error instanceof Error ? error.message : '加载失败');
    } finally {
      if (alive.current) setBusy(false);
    }
  }
  async function start(single = false) {
    if (busy || loading) return;
    if (mode === 'api' && !client) {
      setMessage(
        WEBSITE_MODE
          ? '请在模型设置中保存 LLM 与 Jev 连接；不会自动转为演示。'
          : '真实模式必须先连接本机模型服务；不会转为 mock。',
      );
      return;
    }
    if (task.current) {
      stop('开始新回合');
      task.current.dispose();
      task.current = undefined;
      sceneDirty.current = true;
    }
    setSnapshot(undefined);
    const generation = ++epoch.current;
    setBusy(true);
    try {
      seedPositions(0, goal);
      let session = liveBridge.current.getSession();
      if (sceneDirty.current || task.current || !liveBridge.current.isTaskScene())
        session = await liveBridge.current.loadScene({ seed: 0, goal: [...goal] });
      if (generation !== epoch.current || !alive.current) return;
      if (!session) throw new Error('尚未加载任务场景');
      const provider = mode === 'api' ? client! : new MockDecisionProvider();
      session.startAgent((write) => {
        const skills = new DeterministicBaseline(session!, 0, write, true, goal);
        try {
          const controller = new AgentTaskController(skills, provider, instruction);
          if (single) controller.resume(true);
          task.current = controller;
          return controller;
        } catch (error) {
          skills.dispose();
          throw error;
        }
      });
      ownSession.current = session;
      sceneDirty.current = false;
      session.setPaused(false);
      setMessage(
        mode === 'mock'
          ? '显式 mock 决策；执行与验收为真实 MuJoCo 物理。'
          : '真实模型模式，可能计费；失败不降级规则。',
      );
      publish();
    } catch (error) {
      if (generation === epoch.current && alive.current)
        setMessage(error instanceof Error ? error.message : '启动失败');
    } finally {
      if (alive.current) setBusy(false);
    }
  }
  function resume(single = false) {
    const current = task.current;
    if (current && !current.terminal && ownSession.current === liveBridge.current.getSession()) {
      current.resume(single);
      publish();
    } else void start(single);
  }
  function changeGoal(axis: number, value: number) {
    stop('目标编辑使旧计划失效');
    sceneDirty.current = true;
    setGoal((previous) => previous.map((v, i) => (i === axis ? value : v)));
    // Any edit requires a new scene before starting, even if values are subsequently restored.
    task.current?.dispose();
    setMessage('目标已编辑；开始时会重新加载场景与方块，旧计划不可复用。');
  }
  const running = Boolean(
    snapshot && !['succeeded', 'failed', 'cancelled', 'paused'].includes(snapshot.state),
  );
  const locked = busy || loading;
  const receipts = client?.receipts.filter((r) => r.stamp.runId === snapshot?.runId) ?? [];
  const costs = receipts.filter((r) => r.usage.cost !== undefined);
  const tokens = receipts.filter((r) =>
    ['total_tokens', 'input_tokens', 'prompt_tokens'].some((key) => r.usage[key] !== undefined),
  );
  return (
    <div className="space-y-3 text-xs" aria-label="LeKiwi 智能抓放任务">
      <p>
        仅仿真 · MuJoCo 结构化真值（非视觉） · A /
        lekiwi-v1。空旷平地单方块，不承诺障碍导航或任意目标泛化。
      </p>
      <p>任务模式：{mode === 'mock' ? 'MOCK（不调用模型 API）' : '真实 LLM + Jev（可能计费）'}</p>
      <label className="block">
        决策模式
        <select
          aria-label="智能任务决策模式"
          disabled={locked}
          value={mode}
          onChange={(e) => {
            stop('决策模式改变');
            setMode(e.target.value as 'mock' | 'api');
          }}
          className="w-full bg-surface p-1"
        >
          <option value="mock">显式 mock LLM + Jev</option>
          <option value="api">真实模型服务（无自动回退）</option>
        </select>
      </label>
      <label className="block">
        中文任务指令
        <textarea
          className="w-full rounded border border-border bg-surface p-2"
          rows={3}
          maxLength={2000}
          value={instruction}
          disabled={locked || running}
          onChange={(e) => {
            stop('任务指令修改');
            setInstruction(e.target.value);
          }}
        />
      </label>
      <details open={!WEBSITE_MODE}>
        <summary>高级：目标位置与实验范围</summary>
        <p>语言计划必须遵循安全技能顺序；放置位置以以下场景坐标为准。</p>
        <fieldset disabled={locked || running} className="space-y-1">
          <legend>目标位置（米，暂停后编辑；重新载入回合）</legend>
          {[0, 1].map((axis) => (
            <label key={axis} className="block">
              目标 {axis ? 'Y' : 'X'}
              <input
                className="ml-2 w-24 bg-surface p-1"
                type="number"
                min={TASK_GOAL_LIMITS.min[axis]}
                max={TASK_GOAL_LIMITS.max[axis]}
                step="0.001"
                value={Number.isNaN(goal[axis]) ? '' : goal[axis]}
                onChange={(e) => changeGoal(axis, e.target.valueAsNumber)}
              />
            </label>
          ))}
          <p>Z 固定 {goal[2]}；坐标编辑为受限实验，默认冻结场景已验收。</p>
        </fieldset>
      </details>
      <div className="flex flex-wrap gap-2">
        <Button disabled={locked || running} onClick={() => void start()}>
          开始任务（{mode === 'api' ? '真实 API' : 'mock'}）
        </Button>
        <Button
          disabled={locked || !running}
          onClick={() => {
            task.current?.pause();
            publish();
          }}
        >
          暂停任务
        </Button>
        <Button disabled={locked || snapshot?.state !== 'paused'} onClick={() => resume()}>
          继续任务
        </Button>
        <Button onClick={() => stop()}>停止任务</Button>
      </div>
      <details open={!WEBSITE_MODE}>
        <summary>高级操作与证据导出</summary>
        <div className="mt-2 flex flex-wrap gap-2">
          <Button disabled={locked || running} onClick={() => void load()}>
            加载 A 任务场景
          </Button>
          <Button disabled={locked || running} onClick={() => resume(true)}>
            单技能步进
          </Button>
          <Button disabled={locked} onClick={() => void load()}>
            重置任务场景
          </Button>
          <Button
            disabled={!snapshot}
            onClick={() => {
              const current = task.current;
              if (!current) return;
              downloadBytes(
                new TextEncoder().encode(
                  JSON.stringify(
                    {
                      ...current.snapshot(),
                      results: current.results,
                      failures: current.failures,
                      trace: (current.skills as DeterministicBaseline).trace,
                      receipts,
                    },
                    null,
                    2,
                  ),
                ),
                `lekiwi-agent-${current.snapshot().runId}.json`,
              );
            }}
          >
            导出任务证据
          </Button>
        </div>
      </details>
      <p role="status">{message}</p>
      {snapshot && (
        <div data-testid="agent-task-state" className="space-y-1">
          <p>
            状态：{states[snapshot.state]} · 当前技能：{snapshot.phase}
          </p>
          {snapshot.error && (
            <p role="alert" className="text-warning">
              {snapshot.error}
            </p>
          )}
          <details open={!WEBSITE_MODE}>
            <summary>运行详情与物理证据</summary>
            <p>
              规划 {snapshot.llmCalls}/3 · Jev {snapshot.jevCalls}/60 · 重规划 {snapshot.replans}/2
            </p>
            <p>
              墙钟 {(snapshot.wallMs / 1000).toFixed(1)} s · 仿真 {snapshot.sample.time.toFixed(2)}{' '}
              s · 安全抓持 {snapshot.secure ? '是' : '否'} · 有效搬运{' '}
              {snapshot.transported.toFixed(3)} m
            </p>
            <p>
              双指力 {snapshot.sample.fingerForces.map((v) => v.toFixed(2)).join(' / ')} N ·
              目标支撑 {snapshot.sample.onGoalSupport ? '是' : '否'}
            </p>
            <p>
              Jev：
              {snapshot.decision
                ? `${snapshot.decision.choice} / ${snapshot.decision.grasp} / ${snapshot.decision.diagnosis} / ${snapshot.decision.recovery}`
                : '尚未判定'}
            </p>
            <p>
              已返回费用：
              {costs.length
                ? `$${costs.reduce((sum, r) => sum + r.usage.cost, 0).toFixed(6)}（部分服务报告；非完整账单）`
                : '未报告，不估算为零'}
            </p>
            <p>
              已报告 Token：
              {tokens.length
                ? tokens.reduce(
                    (sum, r) =>
                      sum +
                      (r.usage.total_tokens ??
                        (r.usage.input_tokens ?? r.usage.prompt_tokens ?? 0) +
                          (r.usage.output_tokens ?? r.usage.completion_tokens ?? 0)),
                    0,
                  )
                : '未报告'}{' '}
              · 最近调用耗时{' '}
              {receipts.at(-1)?.elapsedMs === undefined
                ? '未报告'
                : `${receipts.at(-1)!.elapsedMs!.toFixed(0)} ms`}
            </p>
            <p>最近模型：{receipts.at(-1)?.model ?? (mode === 'mock' ? 'mock' : '尚无有效响应')}</p>
          </details>
          {snapshot.plan && (
            <details open>
              <summary>LLM 受约束计划</summary>
              <p>{snapshot.plan.summary}</p>
              <ol className="list-inside list-decimal">
                {snapshot.plan.steps.map((s) => (
                  <li key={s.skill}>
                    {s.skill} · 失败时 {s.onFailure}
                  </li>
                ))}
              </ol>
            </details>
          )}
        </div>
      )}
      <p>
        全局暂停/隐藏/重载会取消任务；任务级暂停可继续。重置前请导出当前证据。
        {!WEBSITE_MODE && 'RL/外控与任务互斥，后台训练作业不会被停止。'}
      </p>
      <details onToggle={(event) => setSettingsOpen(event.currentTarget.open)}>
        <summary>
          {WEBSITE_MODE ? '模型设置 · API_KEY / ChatGPT 订阅' : '模型连接与官方订阅登录'}
        </summary>
        {(!WEBSITE_MODE || settingsOpen) && (
          <ModelConnectionPanel
            onClient={setClient}
            onInvalidate={() => stop('模型配置或账号操作使任务失效')}
          />
        )}
      </details>
    </div>
  );
}