morpheus-cli 0.8.2 → 0.8.3

package/dist/runtime/trinity.js

@@ -210,31 +210,46 @@ ${context ? `CONTEXT FROM ORACLE:\n${context}` : ''}
         const userMessage = new HumanMessage(task);
         const messages = [systemMessage, userMessage];
         try {
+            const startMs = Date.now();
             const response = await this.agent.invoke({ messages });
+            const durationMs = Date.now() - startMs;
             const lastMessage = response.messages[response.messages.length - 1];
             const content = typeof lastMessage.content === 'string'
                 ? lastMessage.content
                 : JSON.stringify(lastMessage.content);
+            const rawUsage = lastMessage.usage_metadata
+                ?? lastMessage.response_metadata?.usage
+                ?? lastMessage.response_metadata?.tokenUsage
+                ?? lastMessage.usage;
+            const inputTokens = rawUsage?.input_tokens ?? 0;
+            const outputTokens = rawUsage?.output_tokens ?? 0;
+            const stepCount = response.messages.filter((m) => m instanceof AIMessage).length;
             const targetSession = sessionId ?? Trinity.currentSessionId ?? 'trinity';
             const history = new SQLiteChatMessageHistory({ sessionId: targetSession });
             try {
                 const persisted = new AIMessage(content);
-                persisted.usage_metadata =
-                    lastMessage.usage_metadata ??
-                        lastMessage.response_metadata?.usage ??
-                        lastMessage.response_metadata?.tokenUsage ??
-                        lastMessage.usage;
-                persisted.provider_metadata = {
-                    provider: trinityConfig.provider,
-                    model: trinityConfig.model,
-                };
+                if (rawUsage)
+                    persisted.usage_metadata = rawUsage;
+                persisted.provider_metadata = { provider: trinityConfig.provider, model: trinityConfig.model };
+                persisted.agent_metadata = { agent: 'trinity' };
+                persisted.duration_ms = durationMs;
                 await history.addMessage(persisted);
             }
             finally {
                 history.close();
             }
             this.display.log('Trinity task completed.', { source: 'Trinity' });
-            return content;
+            return {
+                output: content,
+                usage: {
+                    provider: trinityConfig.provider,
+                    model: trinityConfig.model,
+                    inputTokens,
+                    outputTokens,
+                    durationMs,
+                    stepCount,
+                },
+            };
         }
         catch (err) {
             throw new ProviderError(trinityConfig.provider, err, 'Trinity task execution failed');