solace-agent-mesh 1.4.12__py3-none-any.whl → 1.5.1__py3-none-any.whl

solace_agent_mesh/agent/protocol/event_handlers.py

@@ -2,11 +2,11 @@
 Contains event handling logic for the A2A_ADK_HostComponent.
 """
 
+import logging
 import json
 import asyncio
 from typing import TYPE_CHECKING, Dict, Any
 import fnmatch
-from solace_ai_connector.common.log import log
 from solace_ai_connector.common.message import Message as SolaceMessage
 from ...agent.adk.callbacks import _publish_data_part_status_update
 from ...common.data_parts import ToolResultData
@@ -47,6 +47,7 @@ if TYPE_CHECKING:
     from ..sac.component import SamAgentComponent
 from google.adk.agents.run_config import StreamingMode
 
+log = logging.getLogger(__name__)
 
 def _register_peer_artifacts_in_parent_context(
     parent_task_context: "TaskExecutionContext",
@@ -97,19 +98,6 @@ async def process_event(component, event: Event):
                     component.log_identifier,
                 )
                 return
-            if component.invocation_monitor:
-                component.invocation_monitor.log_message_event(
-                    direction="RECEIVED",
-                    topic=topic,
-                    payload=message.get_payload(),
-                    component_identifier=component.log_identifier,
-                )
-            else:
-                log.warning(
-                    "%s InvocationMonitor not available in component for event on topic %s",
-                    component.log_identifier,
-                    topic,
-                )
             namespace = component.get_config("namespace")
             agent_name = component.get_config("agent_name")
             agent_request_topic = get_agent_request_topic(namespace, agent_name)
@@ -1240,6 +1228,8 @@ async def handle_a2a_response(component, message: SolaceMessage):
                 component.log_identifier,
                 sub_task_id,
             )
+            # Reset the timeout since we received a status update
+            await component.reset_peer_timeout(sub_task_id)
             message.call_acknowledgements()
             return
 

solace_agent_mesh/agent/protocol/protocol_llm.txt

@@ -17,7 +17,7 @@ This is an empty package initialization file and has no public interfaces.
 
 ### event_handlers.py
 **Purpose:** Central hub for processing all events related to the A2A protocol. Routes events to appropriate handlers and manages task lifecycle.
-**Import:** `from solace_agent_mesh.agent.protocol.event_handlers import process_event, handle_a2a_request, handle_agent_card_message, handle_a2a_response, publish_agent_card`
+**Import:** `from solace_agent_mesh.agent.protocol.event_handlers import process_event, handle_a2a_request, handle_agent_card_message, handle_a2a_response, publish_agent_card, handle_sam_event, cleanup_agent_session`
 
 **Functions:**
 - `process_event(component, event: Event) -> None` - Main event router that processes incoming events and delegates to specific handlers based on event type and topic
@@ -25,6 +25,8 @@ This is an empty package initialization file and has no public interfaces.
 - `handle_agent_card_message(component, message: SolaceMessage) -> None` - Processes incoming Agent Card discovery messages and updates peer agent registry
 - `handle_a2a_response(component, message: SolaceMessage) -> None` - Handles responses and status updates from peer agents, manages parallel task completion
 - `publish_agent_card(component) -> None` - Publishes the agent's capabilities and information to the discovery topic
+- `handle_sam_event(component, message: SolaceMessage, topic: str) -> None` - Handles incoming SAM system events like session deletion
+- `cleanup_agent_session(component, session_id: str, user_id: str) -> None` - Cleans up agent-side session data when sessions are deleted
 
 **Internal Helper Functions:**
 - `_register_peer_artifacts_in_parent_context(parent_task_context: "TaskExecutionContext", peer_task_object: Task, log_identifier: str) -> None` - Registers artifacts produced by peer agents in the parent task context
@@ -48,6 +50,16 @@ publish_agent_card(component)
 from solace_agent_mesh.agent.protocol.event_handlers import handle_a2a_request
 
 await handle_a2a_request(component, solace_message)
+
+# Handle SAM system events
+from solace_agent_mesh.agent.protocol.event_handlers import handle_sam_event
+
+handle_sam_event(component, message, topic)
+
+# Clean up session data
+from solace_agent_mesh.agent.protocol.event_handlers import cleanup_agent_session
+
+await cleanup_agent_session(component, "session_123", "user_456")
 ```
 
 **Key Event Flow:**
@@ -56,6 +68,7 @@ await handle_a2a_request(component, solace_message)
    - Agent request topics → `handle_a2a_request()`
    - Discovery topics → `handle_agent_card_message()`
    - Response/status topics → `handle_a2a_response()`
+   - SAM events topics → `handle_sam_event()`
 3. For TIMER events, handles periodic agent card publishing
 4. For CACHE_EXPIRY events, delegates to component's cache handling
 
@@ -65,4 +78,4 @@ await handle_a2a_request(component, solace_message)
 - Integrates with Google ADK for task execution
 - Manages task contexts through `TaskExecutionContext`
 
-# content_hash: 00fc98ca0832801776639518be8b53197c211ac407f79f9a40aa0dc35bf6b447
+# content_hash: 4b8a426320c6f06c621b469c79087c315f5590ef14497ccf17de02eeb66db91d

solace_agent_mesh/agent/protocol/protocol_llm_detail.txt

@@ -0,0 +1,92 @@
+# LLM Summary Detail File
+
+This file is a concatenation of all individual *llm.txt files found in the 'protocol' directory tree. Each section below corresponds to a specific directory's summary file.
+
+================================================================================
+
+## Section 1: solace_agent_mesh/agent/protocol/protocol_llm.txt
+
+**Source file:** `solace_agent_mesh/agent/protocol/protocol_llm.txt`
+
+# DEVELOPER GUIDE: protocol
+
+## Quick Summary
+The `protocol` directory implements the core logic for Agent-to-Agent (A2A) communication. It handles receiving and processing requests, responses, and discovery messages (Agent Cards) over the Solace event mesh. It acts as the bridge between the A2A protocol and the underlying Google ADK execution environment.
+
+## Files Overview
+- `__init__.py` - Empty package initialization file
+- `event_handlers.py` - Contains the primary logic for handling all A2A protocol events, including routing incoming messages, managing task execution, and handling agent discovery
+
+## Developer API Reference
+
+### __init__.py
+**Purpose:** Standard Python package initialization file
+**Import:** `from solace_agent_mesh.agent.protocol import *`
+
+This is an empty package initialization file and has no public interfaces.
+
+### event_handlers.py
+**Purpose:** Central hub for processing all events related to the A2A protocol. Routes events to appropriate handlers and manages task lifecycle.
+**Import:** `from solace_agent_mesh.agent.protocol.event_handlers import process_event, handle_a2a_request, handle_agent_card_message, handle_a2a_response, publish_agent_card, handle_sam_event, cleanup_agent_session`
+
+**Functions:**
+- `process_event(component, event: Event) -> None` - Main event router that processes incoming events and delegates to specific handlers based on event type and topic
+- `handle_a2a_request(component, message: SolaceMessage) -> None` - Handles incoming A2A request messages, starts ADK runner for SendTask requests, and processes CancelTask requests
+- `handle_agent_card_message(component, message: SolaceMessage) -> None` - Processes incoming Agent Card discovery messages and updates peer agent registry
+- `handle_a2a_response(component, message: SolaceMessage) -> None` - Handles responses and status updates from peer agents, manages parallel task completion
+- `publish_agent_card(component) -> None` - Publishes the agent's capabilities and information to the discovery topic
+- `handle_sam_event(component, message: SolaceMessage, topic: str) -> None` - Handles incoming SAM system events like session deletion
+- `cleanup_agent_session(component, session_id: str, user_id: str) -> None` - Cleans up agent-side session data when sessions are deleted
+
+**Internal Helper Functions:**
+- `_register_peer_artifacts_in_parent_context(parent_task_context: "TaskExecutionContext", peer_task_object: Task, log_identifier: str) -> None` - Registers artifacts produced by peer agents in the parent task context
+- `_publish_peer_tool_result_notification(component: "SamAgentComponent", correlation_data: Dict[str, Any], payload_to_queue: Any, log_identifier: str) -> None` - Publishes a ToolResultData status update for a completed peer tool call
+
+**Usage Examples:**
+```python
+# Main event processing (typically called by the SAC framework)
+from solace_agent_mesh.agent.protocol.event_handlers import process_event
+from solace_ai_connector.common.event import Event, EventType
+
+# Process an incoming event
+await process_event(component, event)
+
+# Publish agent discovery card
+from solace_agent_mesh.agent.protocol.event_handlers import publish_agent_card
+
+publish_agent_card(component)
+
+# Handle specific message types (usually called internally by process_event)
+from solace_agent_mesh.agent.protocol.event_handlers import handle_a2a_request
+
+await handle_a2a_request(component, solace_message)
+
+# Handle SAM system events
+from solace_agent_mesh.agent.protocol.event_handlers import handle_sam_event
+
+handle_sam_event(component, message, topic)
+
+# Clean up session data
+from solace_agent_mesh.agent.protocol.event_handlers import cleanup_agent_session
+
+await cleanup_agent_session(component, "session_123", "user_456")
+```
+
+**Key Event Flow:**
+1. `process_event()` receives all events and routes based on type (MESSAGE, TIMER, CACHE_EXPIRY)
+2. For MESSAGE events, routes to specific handlers based on topic patterns:
+   - Agent request topics → `handle_a2a_request()`
+   - Discovery topics → `handle_agent_card_message()`
+   - Response/status topics → `handle_a2a_response()`
+   - SAM events topics → `handle_sam_event()`
+3. For TIMER events, handles periodic agent card publishing
+4. For CACHE_EXPIRY events, delegates to component's cache handling
+
+**Dependencies:**
+- Requires `SamAgentComponent` instance with proper configuration
+- Uses A2A protocol types from `a2a.types`
+- Integrates with Google ADK for task execution
+- Manages task contexts through `TaskExecutionContext`
+
+================================================================================
+

solace_agent_mesh/agent/sac/app.py

@@ -1,3 +1,4 @@
+import logging
 import sys
 import os
 
@@ -7,7 +8,6 @@ sys.path.insert(
 
 from common.utils.asyncio_macos_fix import ensure_asyncio_compatibility
 from .patch_adk import patch_adk
-from solace_ai_connector.common.log import log
 
 ensure_asyncio_compatibility()
 patch_adk()
@@ -29,6 +29,8 @@ from ...agent.utils.artifact_helpers import DEFAULT_SCHEMA_MAX_KEYS
 from ...common.utils.pydantic_utils import SamConfigBase
 from ..tools.tool_config_types import AnyToolConfig
 
+log = logging.getLogger(__name__)
+
 info = {
     "class_name": "SamAgentApp",
     "description": "Custom App class for SAM Agent Host with namespace prefixing and automatic subscription generation.",

solace_agent_mesh/agent/sac/component.py

@@ -2,6 +2,7 @@
 Custom Solace AI Connector Component to Host Google ADK Agents via A2A Protocol.
 """
 
+import logging
 from typing import Any, Dict, Optional, Union, Callable, List, Tuple, TYPE_CHECKING
 import asyncio
 import functools
@@ -15,8 +16,8 @@ import json
 from solace_ai_connector.common.message import (
     Message as SolaceMessage,
 )
-from solace_ai_connector.common.log import log
 from solace_ai_connector.common.event import Event, EventType
+
 from solace_ai_connector.common.utils import import_module
 import inspect
 from pydantic import BaseModel, ValidationError
@@ -71,12 +72,13 @@ from ...agent.tools.peer_agent_tool import (
     PeerAgentTool,
     PEER_TOOL_PREFIX,
 )
-from ...agent.adk.invocation_monitor import InvocationMonitor
 from ...common.middleware.registry import MiddlewareRegistry
 from ...common.constants import DEFAULT_COMMUNICATION_TIMEOUT
 from ...agent.tools.registry import tool_registry
 from ...common.sac.sam_component_base import SamComponentBase
 
+log = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from .app import AgentInitCleanupConfig
     from .task_execution_context import TaskExecutionContext
@@ -248,7 +250,6 @@ class SamAgentComponent(SamComponentBase):
         self._agent_system_instruction_callback: Optional[
             Callable[[CallbackContext, LlmRequest], Optional[str]]
         ] = None
-        self.invocation_monitor: Optional[InvocationMonitor] = None
         self._active_background_tasks = set()
         try:
             self.agent_specific_state: Dict[str, Any] = {}
@@ -389,15 +390,6 @@ class SamAgentComponent(SamComponentBase):
                         raise RuntimeError(
                             f"Agent custom initialization failed: {e}"
                         ) from e
-            try:
-                self.invocation_monitor = InvocationMonitor()
-            except Exception as im_e:
-                log.error(
-                    "%s Failed to initialize InvocationMonitor: %s",
-                    self.log_identifier,
-                    im_e,
-                )
-                self.invocation_monitor = None
 
             # Async init is now handled by the base class `run` method.
             # We still need a future to signal completion from the async thread.
@@ -619,6 +611,42 @@ class SamAgentComponent(SamComponentBase):
             log.warning("%s Failed to claim; it was already completed.", log_id)
             return None
 
+    async def reset_peer_timeout(self, sub_task_id: str):
+        """
+        Resets the timeout for a given peer sub-task.
+        """
+        log_id = f"{self.log_identifier}[ResetTimeout:{sub_task_id}]"
+        log.debug("%s Resetting timeout for peer sub-task.", log_id)
+
+        # Get the original logical task ID from the cache without removing it
+        logical_task_id = self.cache_service.get_data(sub_task_id)
+        if not logical_task_id:
+            log.warning(
+                "%s No active task found for sub-task %s. Cannot reset timeout.",
+                log_id,
+                sub_task_id,
+            )
+            return
+
+        # Get the configured timeout
+        timeout_sec = self.inter_agent_communication_config.get(
+            "request_timeout_seconds", DEFAULT_COMMUNICATION_TIMEOUT
+        )
+
+        # Update the cache with a new expiry
+        self.cache_service.add_data(
+            key=sub_task_id,
+            value=logical_task_id,
+            expiry=timeout_sec,
+            component=self,
+        )
+        log.info(
+            "%s Timeout for sub-task %s has been reset to %d seconds.",
+            log_id,
+            sub_task_id,
+            timeout_sec,
+        )
+
     async def _retrigger_agent_with_peer_responses(
         self,
         results_to_inject: list,
@@ -2050,6 +2078,21 @@ class SamAgentComponent(SamComponentBase):
                     self.log_identifier,
                     len(task_context.produced_artifacts),
                 )
+            
+            # Add token usage summary
+            if task_context:
+                token_summary = task_context.get_token_usage_summary()
+                if token_summary["total_tokens"] > 0:
+                    final_task_metadata["token_usage"] = token_summary
+                    log.info(
+                        "%s Task %s used %d total tokens (input: %d, output: %d, cached: %d)",
+                        self.log_identifier,
+                        logical_task_id,
+                        token_summary["total_tokens"],
+                        token_summary["total_input_tokens"],
+                        token_summary["total_output_tokens"],
+                        token_summary["total_cached_input_tokens"],
+                    )
 
             final_task = a2a.create_final_task(
                 task_id=logical_task_id,
@@ -2991,16 +3034,6 @@ class SamAgentComponent(SamComponentBase):
                         func_name,
                         e,
                     )
-        if self.invocation_monitor:
-            try:
-                self.invocation_monitor.cleanup()
-            except Exception as im_clean_e:
-                log.error(
-                    "%s Error during InvocationMonitor cleanup: %s",
-                    self.log_identifier,
-                    im_clean_e,
-                )
-
         if self._tool_cleanup_hooks:
             log.info(
                 "%s Executing %d tool cleanup hooks...",

solace_agent_mesh/agent/sac/sac_llm.txt

@@ -141,6 +141,8 @@ patch_adk()
   - `get_and_clear_artifact_signals() -> List[Dict[str, Any]]` - Retrieves and clears artifact signals
   - `set_event_loop(loop: asyncio.AbstractEventLoop) -> None` - Stores event loop reference
   - `get_event_loop() -> Optional[asyncio.AbstractEventLoop]` - Retrieves stored event loop
+  - `record_token_usage(input_tokens: int, output_tokens: int, model: str, source: str = "agent", tool_name: Optional[str] = None, cached_input_tokens: int = 0) -> None` - Records token usage for LLM calls
+  - `get_token_usage_summary() -> Dict[str, Any]` - Returns summary of all token usage for the task
 
 **Usage Examples:**
 ```python
@@ -170,6 +172,18 @@ task_context.register_peer_sub_task("sub-task-456", correlation_data)
 completed_data = task_context.claim_sub_task_completion("sub-task-456")
 if completed_data:
     print(f"Sub-task completed: {completed_data}")
+
+# Track token usage
+task_context.record_token_usage(
+    input_tokens=100,
+    output_tokens=50,
+    model="gemini-1.5-pro",
+    source="agent"
+)
+
+# Get usage summary
+usage = task_context.get_token_usage_summary()
+print(f"Total tokens used: {usage['total_tokens']}")
 ```
 
-# content_hash: 79c043cb89eecb2b6a6f05113ff223c5064861f1a537d8e832ebb2738f5801e8
+# content_hash: 27b0000d250131d578d28ab8c4d06c2e2084d888b47102e9b3b133adbc61a169

solace_agent_mesh/agent/sac/sac_llm_detail.txt

@@ -0,0 +1,200 @@
+# LLM Summary Detail File
+
+This file is a concatenation of all individual *llm.txt files found in the 'sac' directory tree. Each section below corresponds to a specific directory's summary file.
+
+================================================================================
+
+## Section 1: solace_agent_mesh/agent/sac/sac_llm.txt
+
+**Source file:** `solace_agent_mesh/agent/sac/sac_llm.txt`
+
+# DEVELOPER GUIDE for the directory: sac
+
+## Quick Summary
+The `sac` (Solace AI Connector) directory provides the core implementation for hosting a Google ADK (Agent Development Kit) agent within the Solace AI Connector framework. It acts as a bridge, enabling ADK agents to communicate using the A2A (Agent-to-Agent) protocol over Solace messaging. This allows for the creation of distributed, collaborative agent systems where agents can delegate tasks, share information, and work together to solve complex problems.
+
+## Files Overview
+- `__init__.py` - Empty package marker file
+- `app.py` - Custom SAC App class that automatically configures Solace subscriptions and broker settings for A2A communication
+- `component.py` - Main SAC Component that hosts the ADK agent, manages its lifecycle, and handles all A2A protocol messaging
+- `patch_adk.py` - Runtime patches for the Google ADK library to enhance or correct its behavior
+- `task_execution_context.py` - State management class that encapsulates all runtime information for a single, in-flight A2A task
+
+## Developer API Reference
+
+### app.py
+**Purpose:** Provides a custom SAC App class that simplifies the configuration of an A2A agent
+**Import:** `from solace_agent_mesh.agent.sac.app import SamAgentApp`
+
+**Classes:**
+- `SamAgentApp(app_info: Dict[str, Any], **kwargs)` - Custom App class for SAM Agent Host with namespace prefixing and automatic subscription generation
+  - `app_schema: Dict` - Class attribute defining comprehensive configuration schema for agent host validation
+
+**Constants/Variables:**
+- `info: Dict[str, str]` - Metadata dictionary about the SamAgentApp class for SAC framework discovery
+
+**Usage Examples:**
+```python
+# SamAgentApp is typically instantiated by the SAC framework from YAML config
+# Example agent-config.yaml:
+# app:
+#   class_name: solace_agent_mesh.agent.sac.app.SamAgentApp
+#   app_config:
+#     namespace: "my-org/production"
+#     agent_name: "customer-support-agent"
+#     model: "gemini-1.5-pro-latest"
+#     tools:
+#       - tool_type: "builtin"
+#         tool_name: "file_search"
+#     agent_card:
+#       description: "An agent that can answer questions about customer accounts."
+#     agent_card_publishing:
+#       interval_seconds: 60
+#     session_service:
+#       type: "memory"
+```
+
+### component.py
+**Purpose:** Core component that hosts a Google ADK agent and bridges communication to A2A protocol
+**Import:** `from solace_agent_mesh.agent.sac.component import SamAgentComponent`
+
+**Classes:**
+- `SamAgentComponent(**kwargs)` - Solace AI Connector component that hosts a Google ADK agent
+  - `process_event(event: Event) -> None` - Main entry point for all SAC framework events
+  - `handle_timer_event(timer_data: Dict[str, Any]) -> None` - Handles scheduled timer events for agent card publishing
+  - `handle_cache_expiry_event(cache_data: Dict[str, Any]) -> None` - Handles cache expiry events for peer agent timeouts
+  - `finalize_task_success(a2a_context: Dict) -> None` - Async method to finalize successful task completion
+  - `finalize_task_canceled(a2a_context: Dict) -> None` - Finalizes task as CANCELED
+  - `finalize_task_error(exception: Exception, a2a_context: Dict) -> None` - Async method to finalize failed tasks
+  - `cleanup() -> None` - Cleans up resources on component shutdown
+  - `set_agent_specific_state(key: str, value: Any) -> None` - Sets key-value pair in agent state dictionary
+  - `get_agent_specific_state(key: str, default: Optional[Any] = None) -> Any` - Retrieves value from agent state
+  - `get_async_loop() -> Optional[asyncio.AbstractEventLoop]` - Returns dedicated asyncio event loop
+  - `set_agent_system_instruction_string(instruction_string: str) -> None` - Sets static system prompt injection
+  - `set_agent_system_instruction_callback(callback_function: Callable) -> None` - Sets dynamic system prompt callback
+  - `get_gateway_id() -> str` - Returns unique identifier for agent host instance
+  - `submit_a2a_task(target_agent_name: str, a2a_message: A2AMessage, user_id: str, user_config: Dict[str, Any], sub_task_id: str) -> str` - Submits task to peer agent
+  - `get_agent_context() -> Dict[str, Any]` - Returns agent context for middleware interactions
+
+**Constants/Variables:**
+- `info: Dict` - Metadata dictionary for SAC framework
+- `CORRELATION_DATA_PREFIX: str` - Prefix for cache keys when tracking peer requests
+- `HOST_COMPONENT_VERSION: str` - Version string of the host component
+
+**Usage Examples:**
+```python
+# Custom initialization function example
+from solace_agent_mesh.agent.sac.component import SamAgentComponent
+
+def initialize_my_agent(host_component: SamAgentComponent, config: dict):
+    """Custom initialization function for the agent."""
+    # Store database connection in agent state
+    db_connection = create_database_connection(config.get('db_url'))
+    host_component.set_agent_specific_state('db_connection', db_connection)
+    
+    # Set custom system instruction
+    host_component.set_agent_system_instruction_string(
+        "You are a specialized customer service agent with access to our database."
+    )
+
+# Tool accessing agent state
+def my_custom_tool(host_component: SamAgentComponent, query: str) -> str:
+    """Tool that uses stored database connection."""
+    db_connection = host_component.get_agent_specific_state('db_connection')
+    if db_connection:
+        return db_connection.execute_query(query)
+    return "Database not available"
+
+# Scheduling async work from synchronous code
+def schedule_background_task(host_component: SamAgentComponent):
+    """Schedule async work on the component's event loop."""
+    loop = host_component.get_async_loop()
+    if loop:
+        asyncio.run_coroutine_threadsafe(my_async_task(), loop)
+```
+
+### patch_adk.py
+**Purpose:** Contains runtime patches for the Google ADK library to enhance behavior
+**Import:** `from solace_agent_mesh.agent.sac.patch_adk import patch_adk`
+
+**Functions:**
+- `patch_adk() -> None` - Applies all necessary patches to the ADK library
+
+**Usage Examples:**
+```python
+from solace_agent_mesh.agent.sac.patch_adk import patch_adk
+
+# Apply patches before using ADK
+patch_adk()
+```
+
+### task_execution_context.py
+**Purpose:** State management class for single, in-flight agent tasks
+**Import:** `from solace_agent_mesh.agent.sac.task_execution_context import TaskExecutionContext`
+
+**Classes:**
+- `TaskExecutionContext(task_id: str, a2a_context: Dict[str, Any])` - Encapsulates runtime state for a single agent task
+  - `cancel() -> None` - Signals that the task should be cancelled
+  - `is_cancelled() -> bool` - Checks if cancellation event has been set
+  - `append_to_streaming_buffer(text: str) -> None` - Appends text to streaming buffer
+  - `flush_streaming_buffer() -> str` - Returns and clears streaming buffer content
+  - `get_streaming_buffer_content() -> str` - Returns buffer content without clearing
+  - `append_to_run_based_buffer(text: str) -> None` - Appends text to run-based response buffer
+  - `register_peer_sub_task(sub_task_id: str, correlation_data: Dict[str, Any]) -> None` - Adds peer sub-task tracking
+  - `claim_sub_task_completion(sub_task_id: str) -> Optional[Dict[str, Any]]` - Atomically retrieves and removes sub-task data
+  - `register_parallel_call_sent(invocation_id: str) -> None` - Registers new parallel tool call
+  - `handle_peer_timeout(sub_task_id: str, correlation_data: Dict, timeout_sec: int, invocation_id: str) -> bool` - Handles peer timeout
+  - `record_parallel_result(result: Dict, invocation_id: str) -> bool` - Records parallel tool call result
+  - `clear_parallel_invocation_state(invocation_id: str) -> None` - Removes completed invocation state
+  - `register_produced_artifact(filename: str, version: int) -> None` - Tracks newly created artifacts
+  - `add_artifact_signal(signal: Dict[str, Any]) -> None` - Adds artifact return signal
+  - `get_and_clear_artifact_signals() -> List[Dict[str, Any]]` - Retrieves and clears artifact signals
+  - `set_event_loop(loop: asyncio.AbstractEventLoop) -> None` - Stores event loop reference
+  - `get_event_loop() -> Optional[asyncio.AbstractEventLoop]` - Retrieves stored event loop
+  - `record_token_usage(input_tokens: int, output_tokens: int, model: str, source: str = "agent", tool_name: Optional[str] = None, cached_input_tokens: int = 0) -> None` - Records token usage for LLM calls
+  - `get_token_usage_summary() -> Dict[str, Any]` - Returns summary of all token usage for the task
+
+**Usage Examples:**
+```python
+from solace_agent_mesh.agent.sac.task_execution_context import TaskExecutionContext
+
+# Create task context
+a2a_context = {
+    "logical_task_id": "task-123",
+    "user_id": "user-456",
+    "session_id": "session-789"
+}
+task_context = TaskExecutionContext("task-123", a2a_context)
+
+# Use streaming buffer
+task_context.append_to_streaming_buffer("Hello ")
+task_context.append_to_streaming_buffer("world!")
+content = task_context.flush_streaming_buffer()  # Returns "Hello world!"
+
+# Track peer sub-tasks
+correlation_data = {
+    "peer_agent_name": "math-agent",
+    "adk_function_call_id": "call-123"
+}
+task_context.register_peer_sub_task("sub-task-456", correlation_data)
+
+# Handle completion
+completed_data = task_context.claim_sub_task_completion("sub-task-456")
+if completed_data:
+    print(f"Sub-task completed: {completed_data}")
+
+# Track token usage
+task_context.record_token_usage(
+    input_tokens=100,
+    output_tokens=50,
+    model="gemini-1.5-pro",
+    source="agent"
+)
+
+# Get usage summary
+usage = task_context.get_token_usage_summary()
+print(f"Total tokens used: {usage['total_tokens']}")
+```
+
+================================================================================
+

solace_agent_mesh/agent/sac/task_execution_context.py

@@ -33,6 +33,13 @@ class TaskExecutionContext:
         self.artifact_signals_to_return: List[Dict[str, Any]] = []
         self.event_loop: Optional[asyncio.AbstractEventLoop] = None
         self.lock: threading.Lock = threading.Lock()
+        
+        # Token usage tracking
+        self.total_input_tokens: int = 0
+        self.total_output_tokens: int = 0
+        self.total_cached_input_tokens: int = 0
+        self.token_usage_by_model: Dict[str, Dict[str, int]] = {}
+        self.token_usage_by_source: Dict[str, Dict[str, int]] = {}
 
     def cancel(self) -> None:
         """Signals that the task should be cancelled."""
@@ -183,3 +190,69 @@ class TaskExecutionContext:
         """Retrieves the stored event loop."""
         with self.lock:
             return self.event_loop
+
+    def record_token_usage(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        model: str,
+        source: str = "agent",
+        tool_name: Optional[str] = None,
+        cached_input_tokens: int = 0,
+    ) -> None:
+        """
+        Records token usage for an LLM call.
+        
+        Args:
+            input_tokens: Number of input/prompt tokens.
+            output_tokens: Number of output/completion tokens.
+            model: Model identifier used for this call.
+            source: Source of the LLM call ("agent" or "tool").
+            tool_name: Tool name if source is "tool".
+            cached_input_tokens: Number of cached input tokens (optional).
+        """
+        with self.lock:
+            # Update totals
+            self.total_input_tokens += input_tokens
+            self.total_output_tokens += output_tokens
+            self.total_cached_input_tokens += cached_input_tokens
+            
+            # Track by model
+            if model not in self.token_usage_by_model:
+                self.token_usage_by_model[model] = {
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "cached_input_tokens": 0,
+                }
+            self.token_usage_by_model[model]["input_tokens"] += input_tokens
+            self.token_usage_by_model[model]["output_tokens"] += output_tokens
+            self.token_usage_by_model[model]["cached_input_tokens"] += cached_input_tokens
+            
+            # Track by source
+            source_key = f"{source}:{tool_name}" if tool_name else source
+            if source_key not in self.token_usage_by_source:
+                self.token_usage_by_source[source_key] = {
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "cached_input_tokens": 0,
+                }
+            self.token_usage_by_source[source_key]["input_tokens"] += input_tokens
+            self.token_usage_by_source[source_key]["output_tokens"] += output_tokens
+            self.token_usage_by_source[source_key]["cached_input_tokens"] += cached_input_tokens
+
+    def get_token_usage_summary(self) -> Dict[str, Any]:
+        """
+        Returns a summary of all token usage for this task.
+        
+        Returns:
+            Dictionary containing total token counts and breakdowns by model and source.
+        """
+        with self.lock:
+            return {
+                "total_input_tokens": self.total_input_tokens,
+                "total_output_tokens": self.total_output_tokens,
+                "total_cached_input_tokens": self.total_cached_input_tokens,
+                "total_tokens": self.total_input_tokens + self.total_output_tokens,
+                "by_model": dict(self.token_usage_by_model),
+                "by_source": dict(self.token_usage_by_source),
+            }