From ee1722fcdfc5deedae8fc777ec63a30b9955c2d0 Mon Sep 17 00:00:00 2001
From: gujishh <baiaoshh@163.com>
Date: Wed, 8 Jul 2026 10:53:54 +0900
Subject: [PATCH] fix(db): record trajectory execution time

---
 tests/db/test_trajectory.py | 28 +++++++++++++++++++---------
 utu/agents/common.py        |  2 ++
 utu/db/trajectory_model.py  |  8 +++++++-
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/tests/db/test_trajectory.py b/tests/db/test_trajectory.py
index 03917f1a..bb2a76d7 100644
--- a/tests/db/test_trajectory.py
+++ b/tests/db/test_trajectory.py
@@ -1,11 +1,21 @@
-# from utu.agents.common import TaskRecorder
-from utu.agents import SimpleAgent
-from utu.db import DBService, TrajectoryModel
+from types import SimpleNamespace
 
+from utu.db import TrajectoryModel
 
-async def test_traj_model():
-    """Test TrajectoryModel. The recorded trajectory should be saved to db and can be visualized."""
-    agent = SimpleAgent(config="simple/base")
-    task_recorder = await agent.run("hello")
-    trajectory = TrajectoryModel.from_task_recorder(task_recorder)
-    DBService.add(trajectory)
+
+def test_from_task_recorder_records_elapsed_time(monkeypatch):
+    monkeypatch.setattr("utu.db.trajectory_model.time.time", lambda: 15.0)
+    recorder = SimpleNamespace(
+        trace_id="trace-1",
+        task="hello",
+        input="",
+        final_output="done",
+        trajectories=[],
+        started_at=10.0,
+    )
+
+    trajectory = TrajectoryModel.from_task_recorder(recorder)
+
+    assert trajectory.time_cost == 5.0
+    assert trajectory.d_input == "hello"
+    assert trajectory.d_output == "done"
diff --git a/utu/agents/common.py b/utu/agents/common.py
index 9f816729..68c0a86b 100644
--- a/utu/agents/common.py
+++ b/utu/agents/common.py
@@ -1,4 +1,5 @@
 import asyncio
+import time
 import traceback
 from collections.abc import AsyncIterator
 from dataclasses import asdict, dataclass, field
@@ -103,6 +104,7 @@ class TaskRecorder(DataClassWithStreamEvents):
 
     # additional infos
     additional_infos: dict = field(default_factory=dict)
+    started_at: float = field(default_factory=time.time)
 
     def to_input_list(self) -> list[TResponseInputItem]:
         return self.get_run_result().to_input_list()
diff --git a/utu/db/trajectory_model.py b/utu/db/trajectory_model.py
index cd34ccb6..3c1a05ea 100644
--- a/utu/db/trajectory_model.py
+++ b/utu/db/trajectory_model.py
@@ -1,4 +1,5 @@
 import json
+import time
 from typing import TYPE_CHECKING
 
 from sqlmodel import Field, SQLModel
@@ -23,11 +24,16 @@ class TrajectoryModel(SQLModel, table=True):
     def from_task_recorder(cls, task_recorder: "TaskRecorder") -> "TrajectoryModel":
         # if isinstance(task_recorder, TaskRecorder):
         d_input = getattr(task_recorder, "task", "") or getattr(task_recorder, "input", "")
+        started_at = getattr(task_recorder, "started_at", None)
+        time_cost = None
+        if isinstance(started_at, (int, float)):
+            time_cost = max(0.0, time.time() - started_at)
+
         return cls(
             trace_id=task_recorder.trace_id,
             trace_url="",
             d_input=d_input,
             d_output=task_recorder.final_output,
             trajectories=json.dumps(task_recorder.trajectories, ensure_ascii=False),
-            time_cost=-1,
+            time_cost=time_cost,
         )