Update sub-mode runtime model handling

agent_templates/agents/a2a_litellm_agent/a2a_server.py

@@ -18,7 +18,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
 import structlog
 
-from agent import LiteLLMAgent
+from agent import LiteLLMAgent, ModelRequestError
 from config import get_config, AgentConfig, A2AConfig
 
 try:
@@ -94,6 +94,7 @@ class A2ATask(BaseModel):
     contextId: str = Field(default_factory=lambda: uuid.uuid4().hex)
     status: A2ATaskStatus
     artifacts: Optional[list[A2AArtifact]] = None
+    metadata: Optional[Dict[str, Any]] = None
 
 
 class A2AResponse(BaseModel):
@@ -392,12 +393,12 @@ class A2AAgentServer:
                     request_id=task_id
                 ) as ctx:
                     ctx.add_tool("a2a_chat")
-                    response_text = await agent.chat(
+                    response = await agent.chat_result(
                         message=user_text,
                         conversation_id=context_id
                     )
             else:
-                response_text = await agent.chat(
+                response = await agent.chat_result(
                     message=user_text,
                     conversation_id=context_id
                 )
@@ -411,9 +412,17 @@ class A2AAgentServer:
             task.artifacts = [
                 A2AArtifact(
                     name="response",
-                    parts=[A2APart(kind="text", text=response_text)]
+                    parts=[A2APart(kind="text", text=response.get("content", ""))]
                 )
             ]
+            task.metadata = {
+                "newapi_request_id": response.get("request_id"),
+                "response_id": response.get("response_id"),
+                "model": response.get("model"),
+                "api_format": response.get("api_format"),
+                "endpoint": response.get("endpoint"),
+                "usage": response.get("usage") or {},
+            }
             self.tasks[task_id] = task
             
             return JSONResponse({
@@ -425,6 +434,9 @@ class A2AAgentServer:
         except Exception as e:
             logger.error("处理消息失败", error=str(e))
             task.status = A2ATaskStatus(state="failed", message=str(e))
+            error_data = {}
+            if isinstance(e, ModelRequestError):
+                error_data = e.to_dict()
             self.tasks[task_id] = task
             
             return JSONResponse({
@@ -432,7 +444,8 @@ class A2AAgentServer:
                 "id": request.id,
                 "error": {
                     "code": -32000,
-                    "message": f"Agent error: {str(e)}"
+                    "message": f"Agent error: {str(e)}",
+                    "data": error_data,
                 }
             })
     

agent_templates/agents/a2a_litellm_agent/agent.py

@@ -6,7 +6,7 @@ LiteLLM Agent 核心模块
 import asyncio
 import json
 import uuid
-from typing import AsyncGenerator, Optional, Dict, Any, List
+from typing import AsyncGenerator, Optional, Dict, Any, List, Union
 from dataclasses import dataclass, field
 from datetime import datetime
 
@@ -47,6 +47,54 @@ class Conversation:
         return [{"role": m.role, "content": m.content} for m in self.messages]
 
 
+@dataclass
+class ModelResult:
+    """Normalized model response metadata for Runtime accounting."""
+
+    content: str
+    usage: Dict[str, int] = field(default_factory=dict)
+    request_id: Optional[str] = None
+    response_id: Optional[str] = None
+    model: Optional[str] = None
+    api_format: str = "openai_chat"
+    endpoint: Optional[str] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "content": self.content,
+            "usage": self.usage,
+            "request_id": self.request_id,
+            "response_id": self.response_id,
+            "model": self.model,
+            "api_format": self.api_format,
+            "endpoint": self.endpoint,
+        }
+
+
+class ModelRequestError(RuntimeError):
+    """Model gateway error carrying request metadata for Runtime logs."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        request_id: Optional[str] = None,
+        status_code: Optional[int] = None,
+        response_text: Optional[str] = None,
+    ):
+        super().__init__(message)
+        self.request_id = request_id
+        self.status_code = status_code
+        self.response_text = response_text
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "request_id": self.request_id,
+            "status_code": self.status_code,
+            "response_text": self.response_text,
+        }
+
+
 class LiteLLMAgent:
     """
     基于LiteLLM的Agent实现
@@ -110,6 +158,8 @@ class LiteLLMAgent:
                 timeout=httpx.Timeout(self.llm_config.timeout),
                 headers={
                     "Authorization": f"Bearer {self.llm_config.api_key}",
+                    "x-api-key": self.llm_config.api_key,
+                    "anthropic-version": "2023-06-01",
                     "Content-Type": "application/json"
                 }
             )
@@ -144,7 +194,7 @@ class LiteLLMAgent:
         message: str,
         conversation_id: Optional[str] = None,
         stream: bool = False
-    ) -> str | AsyncGenerator[str, None]:
+    ) -> Union[str, AsyncGenerator[str, None]]:
         """
         发送消息并获取回复
         
@@ -162,12 +212,73 @@ class LiteLLMAgent:
         conversation.add_message("user", message)
         
         if stream:
+            if self.llm_config.api_format == "anthropic_messages":
+                return self._stream_anthropic_messages_text(conversation)
             return self._stream_chat(conversation)
-        else:
+        result = await self.chat_result_for_conversation(conversation)
+        return result.content
+
+    async def chat_result(
+        self,
+        message: str,
+        conversation_id: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Return assistant text plus usage and NewAPI request metadata."""
+        conversation = self.get_or_create_conversation(conversation_id)
+        conversation.add_message("user", message)
+        return (await self.chat_result_for_conversation(conversation)).to_dict()
+
+    async def chat_result_for_conversation(self, conversation: Conversation) -> ModelResult:
+        """Dispatch to the configured model API format."""
+        if self.llm_config.api_format == "anthropic_messages":
+            if self.llm_config.use_stream:
+                return await self._anthropic_messages_stream(conversation)
+            return await self._anthropic_messages(conversation)
+        if self.llm_config.use_stream:
+            return await self._openai_chat_stream_result(conversation)
         return await self._simple_chat(conversation)
     
-    async def _simple_chat(self, conversation: Conversation) -> str:
-        """非流式对话"""
+    def _request_id_from_response(self, response: httpx.Response, body: Optional[Dict[str, Any]] = None) -> Optional[str]:
+        """Extract NewAPI/OpenAI/Anthropic request ID from headers or body."""
+        for name in (
+            "x-request-id",
+            "request-id",
+            "x-newapi-request-id",
+            "x-litellm-request-id",
+            "anthropic-request-id",
+        ):
+            value = response.headers.get(name)
+            if value:
+                return value
+        if body:
+            return body.get("request_id")
+        return None
+
+    def _normalize_usage(self, usage: Optional[Dict[str, Any]]) -> Dict[str, int]:
+        usage = usage or {}
+        prompt_tokens = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0)
+        completion_tokens = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0)
+        if "input_tokens" in usage or "output_tokens" in usage:
+            total_tokens = prompt_tokens + completion_tokens
+        else:
+            total_tokens = int(usage.get("total_tokens") or prompt_tokens + completion_tokens)
+        return {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+        }
+
+    def _raise_gateway_error(self, exc: httpx.HTTPStatusError, body: Optional[Dict[str, Any]] = None) -> None:
+        request_id = self._request_id_from_response(exc.response, body)
+        raise ModelRequestError(
+            f"Server error '{exc.response.status_code} {exc.response.reason_phrase}' for url '{exc.request.url}'",
+            request_id=request_id,
+            status_code=exc.response.status_code,
+            response_text=exc.response.text[:2000],
+        ) from exc
+
+    async def _simple_chat(self, conversation: Conversation) -> ModelResult:
+        """非流式 OpenAI chat completions 对话"""
         client = await self._get_client()
         
         request_body = {
@@ -184,7 +295,10 @@ class LiteLLMAgent:
                 self.llm_config.chat_endpoint,
                 json=request_body
             )
+            try:
                 response.raise_for_status()
+            except httpx.HTTPStatusError as exc:
+                self._raise_gateway_error(exc)
             
             result = response.json()
             assistant_message = result["choices"][0]["message"]["content"]
@@ -192,8 +306,17 @@ class LiteLLMAgent:
             # 保存助手回复到对话
             conversation.add_message("assistant", assistant_message)
             
-            logger.info("收到回复", length=len(assistant_message))
-            return assistant_message
+            request_id = self._request_id_from_response(response, result)
+            logger.info("收到回复", length=len(assistant_message), request_id=request_id)
+            return ModelResult(
+                content=assistant_message,
+                usage=self._normalize_usage(result.get("usage")),
+                request_id=request_id,
+                response_id=result.get("id"),
+                model=result.get("model") or self.llm_config.model,
+                api_format="openai_chat",
+                endpoint=self.llm_config.chat_endpoint,
+            )
             
         except httpx.HTTPStatusError as e:
             logger.error("HTTP错误", status_code=e.response.status_code, detail=e.response.text)
@@ -202,6 +325,215 @@ class LiteLLMAgent:
             logger.error("请求失败", error=str(e))
             raise
 
+    async def _openai_chat_stream_result(self, conversation: Conversation) -> ModelResult:
+        """OpenAI chat completions stream=true, aggregated into one Runtime artifact."""
+        client = await self._get_client()
+        request_body = {
+            "model": self.llm_config.model,
+            "messages": conversation.to_openai_format(),
+            "temperature": self.llm_config.temperature,
+            "max_tokens": self.llm_config.max_tokens,
+            "stream": True,
+            "stream_options": {"include_usage": True},
+        }
+
+        full_response = ""
+        usage: Dict[str, int] = {}
+        response_id: Optional[str] = None
+        request_id: Optional[str] = None
+        try:
+            async with client.stream("POST", self.llm_config.chat_endpoint, json=request_body) as response:
+                try:
+                    response.raise_for_status()
+                except httpx.HTTPStatusError as exc:
+                    self._raise_gateway_error(exc)
+                request_id = self._request_id_from_response(response)
+                async for line in response.aiter_lines():
+                    if not line.startswith("data: "):
+                        continue
+                    data = line[6:]
+                    if data == "[DONE]":
+                        break
+                    try:
+                        chunk = json.loads(data)
+                    except json.JSONDecodeError:
+                        continue
+                    response_id = response_id or chunk.get("id")
+                    request_id = request_id or chunk.get("request_id")
+                    if chunk.get("usage"):
+                        usage = self._normalize_usage(chunk.get("usage"))
+                    choices = chunk.get("choices") or []
+                    if not choices:
+                        continue
+                    delta = choices[0].get("delta", {})
+                    content = delta.get("content", "")
+                    if content:
+                        full_response += content
+
+            conversation.add_message("assistant", full_response)
+            logger.info("收到流式回复", length=len(full_response), request_id=request_id)
+            return ModelResult(
+                content=full_response,
+                usage=usage,
+                request_id=request_id or response_id,
+                response_id=response_id,
+                model=self.llm_config.model,
+                api_format="openai_chat",
+                endpoint=self.llm_config.chat_endpoint,
+            )
+        except Exception as e:
+            logger.error("流式请求失败", error=str(e))
+            raise
+
+    def _anthropic_payload(self, conversation: Conversation, *, stream: bool = False) -> Dict[str, Any]:
+        system_parts: List[str] = []
+        messages: List[Dict[str, str]] = []
+        for message in conversation.messages:
+            if message.role == "system":
+                system_parts.append(message.content)
+            else:
+                role = "assistant" if message.role == "assistant" else "user"
+                messages.append({"role": role, "content": message.content})
+        payload: Dict[str, Any] = {
+            "model": self.llm_config.model,
+            "messages": messages,
+            "max_tokens": self.llm_config.max_tokens,
+            "stream": stream,
+        }
+        if system_parts:
+            payload["system"] = "\n\n".join(system_parts)
+        return payload
+
+    def _anthropic_text(self, body: Dict[str, Any]) -> str:
+        content = body.get("content") or []
+        texts = [
+            part.get("text", "")
+            for part in content
+            if isinstance(part, dict) and part.get("type") == "text"
+        ]
+        return "".join(texts)
+
+    async def _anthropic_messages(self, conversation: Conversation) -> ModelResult:
+        """Anthropic Messages-compatible call for Claude models."""
+        client = await self._get_client()
+        try:
+            response = await client.post(self.llm_config.messages_endpoint, json=self._anthropic_payload(conversation))
+            try:
+                response.raise_for_status()
+            except httpx.HTTPStatusError as exc:
+                self._raise_gateway_error(exc)
+            result = response.json()
+            assistant_message = self._anthropic_text(result)
+            conversation.add_message("assistant", assistant_message)
+            request_id = self._request_id_from_response(response, result)
+            logger.info("收到 Claude Messages 回复", length=len(assistant_message), request_id=request_id)
+            return ModelResult(
+                content=assistant_message,
+                usage=self._normalize_usage(result.get("usage")),
+                request_id=request_id,
+                response_id=result.get("id"),
+                model=result.get("model") or self.llm_config.model,
+                api_format="anthropic_messages",
+                endpoint=self.llm_config.messages_endpoint,
+            )
+        except Exception as e:
+            logger.error("Claude Messages 请求失败", error=str(e))
+            raise
+
+    async def _anthropic_messages_stream(self, conversation: Conversation) -> ModelResult:
+        """Anthropic Messages stream=true, aggregated into one Runtime artifact."""
+        client = await self._get_client()
+        full_response = ""
+        usage: Dict[str, int] = {}
+        response_id: Optional[str] = None
+        request_id: Optional[str] = None
+        try:
+            async with client.stream(
+                "POST",
+                self.llm_config.messages_endpoint,
+                json=self._anthropic_payload(conversation, stream=True),
+            ) as response:
+                try:
+                    response.raise_for_status()
+                except httpx.HTTPStatusError as exc:
+                    self._raise_gateway_error(exc)
+                request_id = self._request_id_from_response(response)
+                async for line in response.aiter_lines():
+                    if not line.startswith("data: "):
+                        continue
+                    data = line[6:]
+                    if data == "[DONE]":
+                        break
+                    try:
+                        event = json.loads(data)
+                    except json.JSONDecodeError:
+                        continue
+                    event_type = event.get("type")
+                    if event_type == "message_start":
+                        message = event.get("message") or {}
+                        response_id = response_id or message.get("id")
+                        usage = self._normalize_usage(message.get("usage"))
+                    elif event_type == "content_block_delta":
+                        delta = event.get("delta") or {}
+                        text = delta.get("text", "")
+                        if text:
+                            full_response += text
+                    elif event_type == "message_delta":
+                        delta_usage = (event.get("usage") or {})
+                        if delta_usage:
+                            usage = self._normalize_usage({**usage, **delta_usage})
+
+            conversation.add_message("assistant", full_response)
+            logger.info("收到 Claude Messages 流式回复", length=len(full_response), request_id=request_id)
+            return ModelResult(
+                content=full_response,
+                usage=usage,
+                request_id=request_id or response_id,
+                response_id=response_id,
+                model=self.llm_config.model,
+                api_format="anthropic_messages",
+                endpoint=self.llm_config.messages_endpoint,
+            )
+        except Exception as e:
+            logger.error("Claude Messages 流式请求失败", error=str(e))
+            raise
+
+    async def _stream_anthropic_messages_text(self, conversation: Conversation) -> AsyncGenerator[str, None]:
+        """Yield text deltas from Anthropic Messages stream for A2A stream clients."""
+        client = await self._get_client()
+        full_response = ""
+        try:
+            async with client.stream(
+                "POST",
+                self.llm_config.messages_endpoint,
+                json=self._anthropic_payload(conversation, stream=True),
+            ) as response:
+                try:
+                    response.raise_for_status()
+                except httpx.HTTPStatusError as exc:
+                    self._raise_gateway_error(exc)
+                async for line in response.aiter_lines():
+                    if not line.startswith("data: "):
+                        continue
+                    data = line[6:]
+                    if data == "[DONE]":
+                        break
+                    try:
+                        event = json.loads(data)
+                    except json.JSONDecodeError:
+                        continue
+                    if event.get("type") != "content_block_delta":
+                        continue
+                    delta = event.get("delta") or {}
+                    text = delta.get("text", "")
+                    if text:
+                        full_response += text
+                        yield text
+            conversation.add_message("assistant", full_response)
+        except Exception as e:
+            logger.error("Claude Messages 文本流失败", error=str(e))
+            raise
+    
     async def _stream_chat(self, conversation: Conversation) -> AsyncGenerator[str, None]:
         """流式对话"""
         client = await self._get_client()
@@ -211,7 +543,8 @@ class LiteLLMAgent:
             "messages": conversation.to_openai_format(),
             "temperature": self.llm_config.temperature,
             "max_tokens": self.llm_config.max_tokens,
-            "stream": True
+            "stream": True,
+            "stream_options": {"include_usage": True},
         }
         
         full_response = ""
@@ -232,7 +565,10 @@ class LiteLLMAgent:
                         
                         try:
                             chunk = json.loads(data)
-                            delta = chunk.get("choices", [{}])[0].get("delta", {})
+                            choices = chunk.get("choices") or []
+                            if not choices:
+                                continue
+                            delta = choices[0].get("delta", {})
                             content = delta.get("content", "")
                             if content:
                                 full_response += content

agent_templates/agents/a2a_litellm_agent/config.py

@@ -18,8 +18,9 @@ class LiteLLMConfig:
     # 基础URL - 用户提供的LiteLLM服务地址
     base_url: str = "https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io"
     
-    # 完整的chat completions端点
+    # 完整的模型端点
     chat_endpoint: str = field(init=False)
+    messages_endpoint: str = field(init=False)
     
     # API密钥 - 优先使用传入的，否则从环境变量获取
     api_key: Optional[str] = None
@@ -39,6 +40,12 @@ class LiteLLMConfig:
     # 最大token数
     max_tokens: int = 4096
 
+    # API格式：openai_chat 或 anthropic_messages
+    api_format: str = "openai_chat"
+
+    # 是否强制使用流式请求聚合完整响应
+    use_stream: bool = False
+    
     def __post_init__(self):
         self.base_url = (
             os.getenv("LITELLM_BASE_URL")
@@ -52,7 +59,35 @@ class LiteLLMConfig:
             self.api_key = os.getenv("LITELLM_API_KEY")
         if self.model is None:
             self.model = os.getenv("MODEL_NAME") or os.getenv("LITELLM_MODEL", "gpt-4")
+        model_name = (self.model or "").lower()
+
+        self.api_format = (
+            os.getenv("LLM_API_FORMAT")
+            or os.getenv("MODEL_API_FORMAT")
+            or ("anthropic_messages" if "claude" in model_name else "openai_chat")
+        ).lower()
+
+        if "gpt-5.4" in model_name:
+            self.timeout = 600
+            self.use_stream = True
+        if "claude" in model_name:
+            self.timeout = 600
+            self.use_stream = True
+
+        if os.getenv("LITELLM_TIMEOUT") or os.getenv("LLM_TIMEOUT"):
+            self.timeout = int(os.getenv("LITELLM_TIMEOUT") or os.getenv("LLM_TIMEOUT"))
+        if os.getenv("LITELLM_MAX_TOKENS") or os.getenv("LLM_MAX_TOKENS"):
+            self.max_tokens = int(os.getenv("LITELLM_MAX_TOKENS") or os.getenv("LLM_MAX_TOKENS"))
+        if os.getenv("LITELLM_STREAM") or os.getenv("LLM_STREAM"):
+            self.use_stream = (os.getenv("LITELLM_STREAM") or os.getenv("LLM_STREAM", "")).lower() in {
+                "1",
+                "true",
+                "yes",
+                "on",
+            }
+
         self.chat_endpoint = f"{self.base_url}/chat/completions"
+        self.messages_endpoint = f"{self.base_url}/messages"
     
     def validate(self) -> bool:
         """验证配置是否完整"""

api/swarm/orchestrator.py

@@ -154,7 +154,10 @@ class SwarmOrchestrator:
             "model": agent.model,
             "capabilities": agent.capabilities or [],
             "system_prompt": agent.system_prompt,
-            "billing_context": (self.swarm.project_context or {}).get("billing_context") or {},
+            "billing_context": {
+                **((self.swarm.project_context or {}).get("billing_context") or {}),
+                "max_tokens": ((self.swarm.project_context or {}).get("budget") or {}).get("max_tokens"),
+            },
         }
         return self.k8s_manager.deploy_swarm_agent(
             self.swarm_id,
@@ -211,6 +214,7 @@ class SwarmOrchestrator:
                 await self._emit_artifacts(failure_artifacts)
                 await self._emit_phase("failed", "Swarm execution failed")
                 await self._emit_status("failed", {"error": str(e)})
+                await self._emit_usage()
             raise
 
     async def _execute_sequential(self) -> Dict[str, Any]:
@@ -459,8 +463,29 @@ class SwarmOrchestrator:
             if isinstance(response, dict) and response.get("error"):
                 error = response.get("error") or {}
                 message_text = error.get("message") if isinstance(error, dict) else str(error)
-                raise RuntimeError(message_text or "A2A agent returned an error response")
+                error_data = error.get("data") if isinstance(error, dict) else {}
+                if isinstance(error_data, dict):
+                    error_message = SwarmMessage(
+                        message_id=str(uuid.uuid4()),
+                        swarm_id=self.swarm_id,
+                        from_agent_id=client.agent_id,
+                        to_agent_id=None,
+                        message_type="error",
+                        content=message_text or "A2A agent returned an error response",
+                        message_metadata={
+                            "newapi_request_id": error_data.get("request_id"),
+                            "model_status_code": error_data.get("status_code"),
+                            "model_response_preview": (error_data.get("response_text") or "")[:1000],
+                        },
+                    )
+                    self.db.add(error_message)
+                    self.db.commit()
+                runtime_error = RuntimeError(message_text or "A2A agent returned an error response")
+                if isinstance(error_data, dict):
+                    setattr(runtime_error, "request_id", error_data.get("request_id"))
+                raise runtime_error
             usage = self._extract_usage(response)
+            model_metadata = self._extract_model_metadata(response)
             if self.swarm and usage["total_tokens"]:
                 self.swarm.tokens_used += usage["total_tokens"]
                 self.db.commit()
@@ -473,6 +498,8 @@ class SwarmOrchestrator:
                     "summary": "Agent task completed",
                     "result_preview": str(response)[:500],
                     "model_usage": usage,
+                    "newapi_request_id": model_metadata.get("newapi_request_id"),
+                    "model_api_format": model_metadata.get("api_format"),
                 },
             )
 
@@ -484,7 +511,10 @@ class SwarmOrchestrator:
                 to_agent_id=None,  # To orchestrator
                 message_type="response",
                 content=str(response),
-                message_metadata={}
+                message_metadata={
+                    "model_usage": usage,
+                    **model_metadata,
+                },
             )
             self.db.add(response_message)
             self.db.commit()
@@ -512,6 +542,8 @@ class SwarmOrchestrator:
                     "status": "completed",
                     "summary": self._response_summary(response),
                     "runtime_deployment_id": self.swarm_id,
+                    "newapi_request_id": model_metadata.get("newapi_request_id"),
+                    "model_usage": usage,
                 },
             )
 
@@ -532,6 +564,7 @@ class SwarmOrchestrator:
                     "tool_name": "agent_task",
                     "summary": "Agent task failed",
                     "error": str(e),
+                    "newapi_request_id": self._request_id_from_error(e),
                 },
             )
             await self._emit(
@@ -543,6 +576,7 @@ class SwarmOrchestrator:
                     "status": "failed",
                     "summary": str(e),
                     "runtime_deployment_id": self.swarm_id,
+                    "newapi_request_id": self._request_id_from_error(e),
                 },
             )
             raise
@@ -676,13 +710,20 @@ class SwarmOrchestrator:
         context = self._callback_context()
         budget = context.get("budget") or {}
         billing_context = context.get("billing_context") or {}
+        usage = self._aggregate_model_usage_from_messages()
+        if self.swarm and usage["total_tokens"] > self.swarm.tokens_used:
+            self.swarm.tokens_used = usage["total_tokens"]
+            self.db.commit()
+        model_tokens = self.swarm.tokens_used if self.swarm else usage["total_tokens"]
+        if not model_tokens:
+            return
         await self._emit(
             "budget.alert",
             payload={
                 "model_id": billing_context.get("default_model_id") or "unknown",
-                "model_tokens": self.swarm.tokens_used if self.swarm else 0,
-                "prompt_tokens": 0,
-                "completion_tokens": 0,
+                "model_tokens": model_tokens,
+                "prompt_tokens": usage["prompt_tokens"],
+                "completion_tokens": usage["completion_tokens"],
                 "model_cost_usd": 0,
                 "runtime_seconds": 0,
                 "cpu_core_seconds": 0,
@@ -701,6 +742,27 @@ class SwarmOrchestrator:
             },
         )
 
+    def _aggregate_model_usage_from_messages(self) -> Dict[str, int]:
+        """Aggregate persisted model usage metadata for Runtime usage callbacks."""
+        totals = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+        messages = (
+            self.db.query(SwarmMessage)
+            .filter(SwarmMessage.swarm_id == self.swarm_id)
+            .all()
+        )
+        for message in messages:
+            usage = (message.message_metadata or {}).get("model_usage") or {}
+            if not isinstance(usage, dict):
+                continue
+            totals["prompt_tokens"] += int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0)
+            totals["completion_tokens"] += int(usage.get("completion_tokens") or usage.get("output_tokens") or 0)
+            totals["total_tokens"] += int(
+                usage.get("total_tokens")
+                or (usage.get("prompt_tokens") or usage.get("input_tokens") or 0)
+                + (usage.get("completion_tokens") or usage.get("output_tokens") or 0)
+            )
+        return totals
+
     async def _emit(
         self,
         event_type: str,
@@ -845,6 +907,42 @@ class SwarmOrchestrator:
             "total_tokens": total_tokens,
         }
 
+    def _extract_model_metadata(self, response: Any) -> Dict[str, Any]:
+        """Extract model request metadata from nested A2A responses."""
+        metadata = self._find_model_metadata(response) or {}
+        request_id = (
+            metadata.get("newapi_request_id")
+            or metadata.get("request_id")
+            or metadata.get("response_id")
+        )
+        return {
+            "newapi_request_id": request_id,
+            "response_id": metadata.get("response_id"),
+            "model": metadata.get("model"),
+            "api_format": metadata.get("api_format"),
+            "endpoint": metadata.get("endpoint"),
+        }
+
+    def _find_model_metadata(self, value: Any) -> Optional[Dict[str, Any]]:
+        """Find nested model metadata emitted by an A2A agent."""
+        if isinstance(value, dict):
+            if any(key in value for key in ("newapi_request_id", "request_id", "response_id", "api_format")):
+                return value
+            for nested in value.values():
+                found = self._find_model_metadata(nested)
+                if found:
+                    return found
+        elif isinstance(value, list):
+            for item in value:
+                found = self._find_model_metadata(item)
+                if found:
+                    return found
+        return None
+
+    def _request_id_from_error(self, error: Exception) -> Optional[str]:
+        """Best-effort extraction for errors raised after model calls."""
+        return getattr(error, "request_id", None)
+
     def _find_usage_dict(self, value: Any) -> Optional[Dict[str, Any]]:
         """Find a nested usage-like dict containing token counters."""
         if isinstance(value, dict):

api/swarm/router.py

@@ -416,9 +416,18 @@ async def get_swarm_logs(swarm_id: str, db: Session = Depends(get_db)):
         if agent.output:
             log_lines.append(f"last_output={agent.output[:500]}")
         if agent_messages:
-            log_lines.extend(
-                f"{message.created_at.isoformat()} {message.message_type}: {message.content[:300]}"
-                for message in agent_messages[-5:]
+            for message in agent_messages[-5:]:
+                metadata = message.message_metadata or {}
+                request_id = metadata.get("newapi_request_id")
+                usage = metadata.get("model_usage") or {}
+                suffix_parts = []
+                if request_id:
+                    suffix_parts.append(f"newapi_request_id={request_id}")
+                if usage.get("total_tokens"):
+                    suffix_parts.append(f"tokens={usage['total_tokens']}")
+                suffix = f" ({', '.join(suffix_parts)})" if suffix_parts else ""
+                log_lines.append(
+                    f"{message.created_at.isoformat()} {message.message_type}: {message.content[:300]}{suffix}"
                 )
         else:
             log_lines.append("no_runtime_messages_recorded")
@@ -429,6 +438,17 @@ async def get_swarm_logs(swarm_id: str, db: Session = Depends(get_db)):
             "namespace": agent.namespace,
             "pod_name": agent.pod_name,
             "logs": "\n".join(log_lines),
+            "messages": [
+                {
+                    "message_id": message.message_id,
+                    "message_type": message.message_type,
+                    "created_at": message.created_at,
+                    "newapi_request_id": (message.message_metadata or {}).get("newapi_request_id"),
+                    "model_usage": (message.message_metadata or {}).get("model_usage"),
+                    "metadata": message.message_metadata or {},
+                }
+                for message in agent_messages[-20:]
+            ],
         })
 
     return logs

k8s_manager.py

@@ -2449,11 +2449,27 @@ echo "Identity volume initialized successfully"
             role = agent_config.get("role", "worker")
             model = agent_config.get("model", "gpt-4")
             billing_context = agent_config.get("billing_context") or {}
+            model_lower = model.lower()
+            api_format = (
+                billing_context.get("api_format")
+                or ("anthropic_messages" if "claude" in model_lower else "openai_chat")
+            )
+            stream_enabled = billing_context.get("stream")
+            if stream_enabled is None:
+                stream_enabled = "gpt-5.4" in model_lower or "claude" in model_lower
+            elif isinstance(stream_enabled, str):
+                stream_enabled = stream_enabled.lower() in {"1", "true", "yes", "on"}
+            timeout_seconds = int(
+                billing_context.get("timeout_sec")
+                or billing_context.get("timeout_seconds")
+                or (600 if stream_enabled else 300)
+            )
+            max_tokens = int(billing_context.get("max_tokens") or 4096)
 
             # Get template image
             # TODO: Load from template database
             image_map = {
-                "a2a_litellm_agent": "agnettaiji.azurecr.io/ai-agents/a2a-litellm-agent:heicode-v2-gpt54-202605300145-arm64",
+                "a2a_litellm_agent": "agnettaiji.azurecr.io/ai-agents/a2a-litellm-agent:heicode-v2-runtime-20260531214518-arm64",
                 "code_manager_agent": "agnettaiji.azurecr.io/ai-agents/code-manager-agent:latest"
             }
             image = image_map.get(template, image_map["a2a_litellm_agent"])
@@ -2466,6 +2482,14 @@ echo "Identity volume initialized successfully"
                 client.V1EnvVar(name="MODEL_NAME", value=model),
                 client.V1EnvVar(name="POD_NAME", value=pod_name),
                 client.V1EnvVar(name="NAMESPACE", value=namespace),
+                client.V1EnvVar(name="MODEL_API_FORMAT", value=api_format),
+                client.V1EnvVar(name="LLM_API_FORMAT", value=api_format),
+                client.V1EnvVar(name="LITELLM_STREAM", value="true" if stream_enabled else "false"),
+                client.V1EnvVar(name="LLM_STREAM", value="true" if stream_enabled else "false"),
+                client.V1EnvVar(name="LITELLM_TIMEOUT", value=str(timeout_seconds)),
+                client.V1EnvVar(name="LLM_TIMEOUT", value=str(timeout_seconds)),
+                client.V1EnvVar(name="LITELLM_MAX_TOKENS", value=str(max_tokens)),
+                client.V1EnvVar(name="LLM_MAX_TOKENS", value=str(max_tokens)),
             ]
 
             gateway_url = (

scripts/e2e_complete_project.py

@@ -0,0 +1,311 @@
+#!/usr/bin/env python3
+"""End-to-end sub-mode Runtime check for a complete generated project.
+
+This script is intentionally outside the normal unit-test suite because it
+requires a live agent-manager Runtime, Kubernetes, and a model gateway.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+from typing import Any
+
+
+DEFAULT_OBJECTIVE = """Return ONLY valid JSON:
+{"files":[{"path":"...","content":"..."}],"run_tests":"python -m unittest discover -s tests -v","smoke_test":"python -m textstats_cli samples/example.txt --json"}.
+
+Build a complete, compact Python stdlib project named textstats_cli.
+Include exactly these files:
+- pyproject.toml
+- README.md
+- textstats_cli/__main__.py
+- textstats_cli/core.py
+- tests/test_core.py
+- samples/example.txt
+
+Features:
+- CLI accepts a text file path.
+- --json outputs JSON.
+- Default output is human readable.
+- Report line count, word count, character count, top 5 words excluding common stopwords, and estimated reading time.
+- Tests must cover counting, stopword filtering, JSON-safe result shape, and missing-file error handling.
+
+Constraints:
+- Use only the Python standard library.
+- Keep the project small enough for one response.
+- No placeholders.
+- No markdown fences.
+- No prose outside the JSON object.
+"""
+
+
+TERMINAL_STATUSES = {"completed", "failed", "stopped"}
+
+
+def http_json(method: str, url: str, payload: dict[str, Any] | None = None, headers: dict[str, str] | None = None) -> Any:
+    data = None if payload is None else json.dumps(payload).encode("utf-8")
+    request_headers = {"Content-Type": "application/json", **(headers or {})}
+    request = urllib.request.Request(url, data=data, headers=request_headers, method=method)
+    try:
+        with urllib.request.urlopen(request, timeout=60) as response:
+            return json.loads(response.read().decode("utf-8"))
+    except urllib.error.HTTPError as exc:
+        body = exc.read().decode("utf-8", errors="replace")
+        raise RuntimeError(f"{method} {url} failed with HTTP {exc.code}: {body}") from exc
+
+
+def http_bytes(url: str) -> bytes:
+    with urllib.request.urlopen(url, timeout=120) as response:
+        return response.read()
+
+
+def extract_json_object(text: str) -> dict[str, Any]:
+    raw = text.strip()
+    fenced = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", raw, re.S)
+    if fenced:
+        raw = fenced.group(1)
+    else:
+        start = raw.find("{")
+        end = raw.rfind("}")
+        if start < 0 or end <= start:
+            raise ValueError("artifact does not contain a JSON object")
+        raw = raw[start : end + 1]
+    return json.loads(raw)
+
+
+def safe_write_project(project_dir: Path, files: list[dict[str, Any]]) -> None:
+    project_root = project_dir.resolve()
+    for item in files:
+        relative_path = item.get("path")
+        content = item.get("content")
+        if not isinstance(relative_path, str) or not relative_path:
+            raise ValueError(f"invalid file path in artifact: {item!r}")
+        if not isinstance(content, str):
+            raise ValueError(f"invalid content for {relative_path}")
+        target = (project_root / relative_path).resolve()
+        if project_root not in target.parents and target != project_root:
+            raise ValueError(f"artifact attempted to write outside project: {relative_path}")
+        target.parent.mkdir(parents=True, exist_ok=True)
+        target.write_text(content, encoding="utf-8")
+
+
+def run_command(command: str, cwd: Path) -> subprocess.CompletedProcess[str]:
+    env = os.environ.copy()
+    shim_dir = None
+    if shutil.which("python") is None:
+        shim_dir = Path(tempfile.mkdtemp(prefix="heicode-python-shim-"))
+        (shim_dir / "python").symlink_to(sys.executable)
+        env["PATH"] = str(shim_dir) + os.pathsep + env.get("PATH", "")
+    return subprocess.run(
+        command,
+        cwd=str(cwd),
+        env=env,
+        shell=True,
+        text=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        timeout=120,
+        check=False,
+    )
+
+
+def build_payload(args: argparse.Namespace) -> dict[str, Any]:
+    return {
+        "orchestration_plan": {
+            "sub_mode": "agile",
+            "objective": args.objective,
+            "user_context": {"user_id": args.user_id},
+            "agents": [
+                {
+                    "role": "backend",
+                    "template": "a2a_litellm_agent",
+                    "model": args.model,
+                    "capabilities": ["code", "test"],
+                }
+            ],
+            "agile_context": {"max_iterations": 1, "stage": "development"},
+            "budget": {"max_duration_sec": args.timeout_seconds, "max_tokens": args.max_tokens},
+            "billing_context": {
+                "provider": "newapi",
+                "default_model_id": args.model,
+                "model_gateway_url": args.model_gateway_url,
+                "api_format": args.api_format,
+                "stream": args.stream,
+                "timeout_sec": args.model_timeout_seconds,
+                "max_tokens": args.max_tokens,
+            },
+            "metadata": {"correlation_id": args.correlation_id},
+        },
+        "callback": {"url": args.callback_url, "method": "POST"},
+    }
+
+
+def poll_swarm(base_url: str, swarm_id: str, timeout_seconds: int, poll_interval: int) -> dict[str, Any]:
+    deadline = time.time() + timeout_seconds
+    last_line = None
+    while time.time() < deadline:
+        status = http_json("GET", f"{base_url}/api/swarms/{swarm_id}")
+        line = {
+            "status": status.get("status"),
+            "phase": status.get("phase"),
+            "progress": status.get("progress"),
+            "artifact_count": len(status.get("artifacts") or []),
+            "tokens_used": (status.get("metrics") or {}).get("tokens_used"),
+        }
+        if line != last_line:
+            print("status:", json.dumps(line, ensure_ascii=False))
+            last_line = line
+        if status.get("status") in TERMINAL_STATUSES:
+            return status
+        time.sleep(poll_interval)
+    raise TimeoutError(f"swarm {swarm_id} did not finish within {timeout_seconds}s")
+
+
+def assert_runtime_observability(base_url: str, swarm_id: str, status: dict[str, Any]) -> None:
+    metrics = status.get("metrics") or {}
+    tokens_used = int(metrics.get("tokens_used") or 0)
+    if tokens_used <= 0:
+        raise AssertionError(f"Runtime tokens_used must be > 0, got {tokens_used}")
+
+    logs = http_json("GET", f"{base_url}/api/swarms/{swarm_id}/logs")
+    request_ids: list[str] = []
+    usage_totals: list[int] = []
+    for agent in logs.get("agents") or []:
+        for message in agent.get("messages") or []:
+            if message.get("newapi_request_id"):
+                request_ids.append(message["newapi_request_id"])
+            usage = message.get("model_usage") or {}
+            if usage.get("total_tokens"):
+                usage_totals.append(int(usage["total_tokens"]))
+    if not request_ids:
+        raise AssertionError("Runtime logs must include at least one NewAPI request_id")
+    if not usage_totals:
+        raise AssertionError("Runtime logs must include model usage with total_tokens")
+    print("observability:", json.dumps({"request_ids": request_ids, "usage_totals": usage_totals}, ensure_ascii=False))
+
+
+def run_project_validation(base_url: str, status: dict[str, Any], output_dir: Path) -> None:
+    artifacts = status.get("artifacts") or []
+    if status.get("status") != "completed":
+        raise AssertionError(f"swarm did not complete: {status.get('error_message')}")
+    if not artifacts:
+        raise AssertionError("completed swarm returned no artifacts")
+
+    artifact = artifacts[0]
+    download_path = (artifact.get("metadata") or {}).get("download_path")
+    if not download_path:
+        raise AssertionError("artifact metadata is missing download_path")
+
+    artifact_bytes = http_bytes(f"{base_url}{download_path}")
+    artifact_hash = "sha256:" + hashlib.sha256(artifact_bytes).hexdigest()
+    expected_hash = (artifact.get("metadata") or {}).get("content_hash")
+    if expected_hash and artifact_hash != expected_hash:
+        raise AssertionError(f"artifact hash mismatch: expected {expected_hash}, got {artifact_hash}")
+
+    artifact_text = artifact_bytes.decode("utf-8")
+    artifact_json = extract_json_object(artifact_text)
+    files = artifact_json.get("files")
+    if not isinstance(files, list) or len(files) < 5:
+        raise AssertionError("artifact must contain a multi-file project")
+
+    project_dir = output_dir / "project"
+    if project_dir.exists():
+        shutil.rmtree(project_dir)
+    project_dir.mkdir(parents=True)
+    safe_write_project(project_dir, files)
+
+    required_paths = {
+        "pyproject.toml",
+        "README.md",
+        "textstats_cli/__main__.py",
+        "textstats_cli/core.py",
+        "tests/test_core.py",
+        "samples/example.txt",
+    }
+    actual_paths = {str(path.relative_to(project_dir)) for path in project_dir.rglob("*") if path.is_file()}
+    missing = sorted(required_paths - actual_paths)
+    if missing:
+        raise AssertionError(f"generated project missing required files: {missing}")
+
+    for label, command in (
+        ("run_tests", artifact_json.get("run_tests")),
+        ("smoke_test", artifact_json.get("smoke_test")),
+    ):
+        if not isinstance(command, str) or not command.strip():
+            raise AssertionError(f"artifact missing {label}")
+        result = run_command(command, project_dir)
+        print(f"{label}: {command}")
+        print(result.stdout)
+        if result.returncode != 0:
+            raise AssertionError(f"{label} failed with exit code {result.returncode}")
+
+    print("artifact:", json.dumps({
+        "artifact_id": artifact.get("artifact_id"),
+        "uri": artifact.get("uri"),
+        "size_bytes": artifact.get("size_bytes"),
+        "content_hash": artifact_hash,
+        "project_dir": str(project_dir),
+    }, ensure_ascii=False))
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--base-url", default=os.getenv("AGENT_MANAGER_URL", "http://127.0.0.1:8000"))
+    parser.add_argument("--model-gateway-url", default=os.getenv("HEICODE_NEWAPI_BASE_URL", "https://code.xinghanlab.com/v1"))
+    parser.add_argument("--model", default=os.getenv("HEICODE_E2E_MODEL", "gpt-5.4"))
+    parser.add_argument("--api-format", default=os.getenv("HEICODE_E2E_API_FORMAT", "openai_chat"))
+    parser.add_argument("--stream", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--max-tokens", type=int, default=int(os.getenv("HEICODE_E2E_MAX_TOKENS", "6000")))
+    parser.add_argument("--timeout-seconds", type=int, default=int(os.getenv("HEICODE_E2E_TIMEOUT_SECONDS", "1200")))
+    parser.add_argument("--model-timeout-seconds", type=int, default=int(os.getenv("HEICODE_E2E_MODEL_TIMEOUT_SECONDS", "600")))
+    parser.add_argument("--poll-interval", type=int, default=10)
+    parser.add_argument("--user-id", default="heicode-complete-project-e2e")
+    parser.add_argument("--callback-url", default="http://127.0.0.1:9/heicode-callback")
+    parser.add_argument("--correlation-id", default=f"heicode-complete-project-e2e-{int(time.time())}")
+    parser.add_argument("--idempotency-key", default=f"complete-project-e2e-{int(time.time())}")
+    parser.add_argument("--objective", default=DEFAULT_OBJECTIVE)
+    parser.add_argument("--output-dir", type=Path, default=None)
+    return parser.parse_args()
+
+
+def main() -> int:
+    args = parse_args()
+    base_url = args.base_url.rstrip("/")
+    output_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="heicode-complete-project-e2e-"))
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    payload = build_payload(args)
+    created = http_json(
+        "POST",
+        f"{base_url}/api/swarms",
+        payload,
+        headers={"X-Idempotency-Key": args.idempotency_key},
+    )
+    swarm_id = created["swarm_id"]
+    print("created:", json.dumps({"swarm_id": swarm_id, "output_dir": str(output_dir)}, ensure_ascii=False))
+
+    status = poll_swarm(base_url, swarm_id, args.timeout_seconds, args.poll_interval)
+    run_project_validation(base_url, status, output_dir)
+    assert_runtime_observability(base_url, swarm_id, status)
+    print("PASS: complete project E2E succeeded")
+    return 0
+
+
+if __name__ == "__main__":
+    try:
+        raise SystemExit(main())
+    except Exception as exc:
+        print(f"FAIL: {exc}", file=sys.stderr)
+        raise SystemExit(1)