diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8c96ca356..29e05db55 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -18,7 +18,7 @@ jobs: lint: timeout-minutes: 10 name: lint - runs-on: ${{ github.repository == 'stainless-sdks/agentex-sdk-python' && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} + runs-on: ${{ startsWith(github.repository, 'stainless-sdks/') && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} if: (github.event_name == 'push' || github.event.pull_request.head.repo.fork) && (github.event_name != 'push' || github.event.head_commit.message != 'codegen metadata') steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -44,7 +44,7 @@ jobs: permissions: contents: read id-token: write - runs-on: ${{ github.repository == 'stainless-sdks/agentex-sdk-python' && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} + runs-on: ${{ startsWith(github.repository, 'stainless-sdks/') && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 @@ -87,7 +87,7 @@ jobs: test: timeout-minutes: 10 name: test - runs-on: ${{ github.repository == 'stainless-sdks/agentex-sdk-python' && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} + runs-on: ${{ startsWith(github.repository, 'stainless-sdks/') && 'depot-ubuntu-24.04' || 'ubuntu-latest' }} if: github.event_name == 'push' || github.event.pull_request.head.repo.fork steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 diff --git a/.stats.yml b/.stats.yml index 2764fd11e..3381aa342 100644 --- a/.stats.yml +++ b/.stats.yml @@ -1,4 +1,4 @@ configured_endpoints: 75 -openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/sgp/agentex-sdk-b2df5f506330ad5fba5a0d518ab8a4bcf876e8c3684a4fe0d0cc5102fd9c569e.yml -openapi_spec_hash: 132e9efdb0535d9594abadf799431cf5 +openapi_spec_url: https://storage.googleapis.com/stainless-sdk-openapi-specs/sgp/agentex-sdk-132fed556e527565b396688a09e873b51ae42c6ad11757d49fc5ef1f2091c943.yml +openapi_spec_hash: 0c1d8518f9f60bb8dc325e26fe1ebadc config_hash: 593e89b291976a5e84e4c3c3f8324354 diff --git a/src/agentex/lib/adk/__init__.py b/src/agentex/lib/adk/__init__.py index 25b858485..d5be0ac52 100644 --- a/src/agentex/lib/adk/__init__.py +++ b/src/agentex/lib/adk/__init__.py @@ -29,6 +29,10 @@ from agentex.lib.adk._modules.tasks import TasksModule from agentex.lib.adk._modules.tracing import TracingModule, TurnSpan +# Data-source refs for lineage (SGP-6513); implementation lives in core.tracing +from agentex.lib.core.tracing import lineage +from agentex.lib.core.tracing.lineage import DataSourceRef, data_sources + # Unified harness surface (AGX1-375) from agentex.lib.core.harness import ( UnifiedEmitter, @@ -67,6 +71,10 @@ "events", "agent_task_tracker", "TurnSpan", + # Lineage data-source refs (SGP-6513) + "lineage", + "DataSourceRef", + "data_sources", # Checkpointing / LangGraph "create_checkpointer", "stream_langgraph_events", diff --git a/src/agentex/lib/adk/providers/_modules/sync_provider.py b/src/agentex/lib/adk/providers/_modules/sync_provider.py index 86696a2b5..120915eec 100644 --- a/src/agentex/lib/adk/providers/_modules/sync_provider.py +++ b/src/agentex/lib/adk/providers/_modules/sync_provider.py @@ -19,6 +19,7 @@ from agentex import AsyncAgentex from agentex.lib.utils.logging import make_logger from agentex.lib.core.tracing.tracer import AsyncTracer +from agentex.lib.core.tracing.lineage import merge_refs_into_data, resolve_refs_from_items logger = make_logger(__name__) @@ -185,6 +186,9 @@ async def get_response( "new_items": new_items, "final_output": final_output, } + lineage_refs = resolve_refs_from_items(new_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return response else: @@ -303,6 +307,9 @@ async def stream_response( "new_items": new_items, "final_output": final_response_text if final_response_text else None, } + lineage_refs = resolve_refs_from_items(new_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) finally: # End the span after all events have been yielded await trace.end_span(span) diff --git a/src/agentex/lib/core/harness/tracer.py b/src/agentex/lib/core/harness/tracer.py index bf37bad30..34cd95616 100644 --- a/src/agentex/lib/core/harness/tracer.py +++ b/src/agentex/lib/core/harness/tracer.py @@ -6,6 +6,17 @@ from agentex.lib.core.harness.types import OpenSpan, CloseSpan, SpanSignal +try: + from agentex.lib.core.tracing.lineage import resolve_refs, merge_refs_into_data +except Exception: # keep the harness importable without optional tracing deps + + def resolve_refs(tool_name: str, arguments: dict[str, Any] | None) -> list[dict[str, Any]]: # noqa: ARG001 + return [] + + def merge_refs_into_data(data: dict[str, Any] | None, refs: list[dict[str, Any]]) -> dict[str, Any]: # noqa: ARG001 + return dict(data or {}) + + try: from agentex.lib.utils.logging import make_logger @@ -80,6 +91,11 @@ async def handle(self, signal: SpanSignal) -> None: task_id=self.task_id, ) if span is not None: + if signal.kind == "tool": + refs = resolve_refs(signal.name, signal.input if isinstance(signal.input, dict) else {}) + if refs: + data = span.data if isinstance(span.data, dict) else {} + span.data = merge_refs_into_data(data, refs) self._open[signal.key] = span elif isinstance(signal, CloseSpan): span = self._open.pop(signal.key, None) diff --git a/src/agentex/lib/core/services/adk/providers/openai.py b/src/agentex/lib/core/services/adk/providers/openai.py index a2513ea01..cc411dc30 100644 --- a/src/agentex/lib/core/services/adk/providers/openai.py +++ b/src/agentex/lib/core/services/adk/providers/openai.py @@ -25,6 +25,7 @@ from agentex.lib.utils.temporal import heartbeat_if_in_workflow from agentex.lib.core.tracing.tracer import AsyncTracer from agentex.lib.core.harness.emitter import UnifiedEmitter +from agentex.lib.core.tracing.lineage import merge_refs_into_data, resolve_refs_from_items from agentex.types.task_message_update import StreamTaskMessageFull from agentex.types.task_message_content import ( TextContent, @@ -286,13 +287,17 @@ async def run_agent( result = await Runner.run(starting_agent=agent, input=input_list) if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return result @@ -431,13 +436,17 @@ async def run_agent_auto_send( result = await Runner.run(starting_agent=agent, input=input_list) if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) tool_call_map: dict[str, Any] = {} @@ -646,13 +655,17 @@ async def run_agent_streamed( result = Runner.run_streamed(starting_agent=agent, input=input_list) if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return result @@ -906,12 +919,16 @@ async def run_agent_streamed_auto_send( raise if span: + serialized_items = [ + item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item + for item in result.new_items + ] span.output = { - "new_items": [ - item.raw_item.model_dump() if isinstance(item.raw_item, BaseModel) else item.raw_item - for item in result.new_items - ], + "new_items": serialized_items, "final_output": result.final_output, } + lineage_refs = resolve_refs_from_items(serialized_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) return result diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py index 7c8690f21..c985d5e65 100644 --- a/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py +++ b/src/agentex/lib/core/temporal/plugins/openai_agents/models/temporal_streaming_model.py @@ -64,6 +64,7 @@ from agentex.lib import adk from agentex.lib.utils.logging import make_logger from agentex.lib.core.tracing.tracer import AsyncTracer +from agentex.lib.core.tracing.lineage import merge_refs_into_data, resolve_refs_from_items from agentex.types.task_message_delta import TextDelta, ToolRequestDelta, ReasoningContentDelta, ReasoningSummaryDelta from agentex.types.task_message_update import StreamTaskMessageFull, StreamTaskMessageDelta from agentex.types.task_message_content import TextContent, ReasoningContent, ToolRequestContent, ToolResponseContent @@ -1257,6 +1258,9 @@ async def get_response( output_data["tool_outputs"] = tool_outputs span.output = output_data + lineage_refs = resolve_refs_from_items(new_items) + if lineage_refs: + span.data = merge_refs_into_data(span.data, lineage_refs) # Streaming-only metrics. Token counters and the success request # counter are emitted by LLMMetricsHooks.on_llm_end so they fire diff --git a/src/agentex/lib/core/tracing/lineage.py b/src/agentex/lib/core/tracing/lineage.py new file mode 100644 index 000000000..75eaffdc0 --- /dev/null +++ b/src/agentex/lib/core/tracing/lineage.py @@ -0,0 +1,174 @@ +"""Data-source reference capture for lineage: tools declare which sources they +touch and the refs land in span data under the ``sgp.lineage.refs`` key.""" + +from __future__ import annotations + +import re +import json +from typing import Any, Literal, Callable, Iterable + +from pydantic import Field, BaseModel, field_validator + +try: + from agentex.lib.utils.logging import make_logger + + logger = make_logger(__name__) +except Exception: # ddtrace may be absent in some envs; fall back to stdlib + import logging + + logger = logging.getLogger(__name__) + +LINEAGE_REFS_KEY = "sgp.lineage.refs" + +# The URI arm of the lineage namespace identifier rule (namespace-conventions.md): +# lowercase scheme and host (dots/hyphens only — normalize `_` to `-`), one optional path segment. +_URI_NAMESPACE_RE = re.compile(r"^[a-z][a-z0-9._-]*://[a-z0-9.-]+(/[a-zA-Z0-9._-]*)?$") + +RefResolver = Callable[[dict[str, Any]], "list[DataSourceRef]"] + + +class DataSourceRef(BaseModel): + """One data source a tool call touched, as a lineage coordinate.""" + + namespace: str = Field(max_length=512) + name: str = Field(min_length=1, max_length=512) + version: str | None = Field(default=None, max_length=256) + role: Literal["input", "output"] = "input" + + def __init__(self, namespace: str | None = None, name: str | None = None, **kwargs: Any) -> None: + if namespace is not None: + kwargs["namespace"] = namespace + if name is not None: + kwargs["name"] = name + super().__init__(**kwargs) + + @field_validator("namespace") + @classmethod + def _namespace_is_uri_form(cls, value: str) -> str: + if not _URI_NAMESPACE_RE.match(value): + raise ValueError(f"namespace must be URI-form (scheme://system), got: {value!r}") + return value + + +class _ToolSources(BaseModel): + refs: list[DataSourceRef] = Field(default_factory=list) + resolver: RefResolver | None = None + + model_config = {"arbitrary_types_allowed": True} + + +_tool_sources: dict[str, _ToolSources] = {} + + +def register_tool_sources( + tool_name: str, + refs: Iterable[DataSourceRef] | None = None, + resolver: RefResolver | None = None, +) -> None: + """Declare the data sources a tool touches, keyed by its tool name. + + Use for tools the agent does not own (e.g. MCP proxy tools). Static refs and + a resolver over the tool's parsed arguments may be combined; repeated + registration for the same name replaces the prior entry. The registry is + process-wide: co-located agents sharing a tool name share (and overwrite) + one entry, so disambiguate shared names before co-locating agent types. + """ + _tool_sources[tool_name] = _ToolSources(refs=list(refs or []), resolver=resolver) + + +def data_sources(*refs: DataSourceRef, resolver: RefResolver | None = None) -> Callable[[Any], Any]: + """Decorator form of ``register_tool_sources`` for tools the agent owns. + + Works below or above ``@function_tool``: the tool name is taken from the + decorated object's ``name`` attribute when present, else ``__name__``. + """ + + def _register(obj: Any) -> Any: + tool_name = getattr(obj, "name", None) or getattr(obj, "__name__", None) + if isinstance(tool_name, str) and tool_name: + register_tool_sources(tool_name, refs=refs, resolver=resolver) + else: + logger.warning("data_sources could not determine a tool name for %r; refs not registered", obj) + return obj + + return _register + + +def clear_tool_sources() -> None: + """Reset the registry (test isolation).""" + _tool_sources.clear() + + +def resolve_refs(tool_name: str, arguments: dict[str, Any] | None) -> list[dict[str, Any]]: + """Resolve registered refs for one tool call to serialized, deduplicated dicts. + + Resolver failures are logged and swallowed: ref capture must never break a + tool call or its tracing. + """ + entry = _tool_sources.get(tool_name) + if entry is None: + return [] + refs = list(entry.refs) + if entry.resolver is not None: + try: + refs.extend(entry.resolver(arguments or {})) + except Exception: + logger.warning("data-source resolver for tool %s failed; static refs kept", tool_name, exc_info=True) + return _dedupe(refs) + + +def resolve_refs_from_items(items: Iterable[Any]) -> list[dict[str, Any]]: + """Resolve refs across serialized run items, matching ``function_call`` entries. + + Accepts the item dicts the providers already build for span output; string + ``arguments`` are parsed as JSON for resolver-based registrations. + """ + refs: list[dict[str, Any]] = [] + for item in items: + if not isinstance(item, dict) or item.get("type") != "function_call": + continue + tool_name = item.get("name") + if not isinstance(tool_name, str) or not tool_name: + continue + arguments = item.get("arguments") + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except (ValueError, TypeError): + arguments = {} + refs.extend(resolve_refs(tool_name, arguments if isinstance(arguments, dict) else {})) + return _dedupe_dicts(refs) + + +def record(span: Any, refs: Iterable[DataSourceRef]) -> None: + """Attach refs to a manually managed span (no-op when the span is None).""" + if span is None: + return + merged = merge_refs_into_data(getattr(span, "data", None), _dedupe(list(refs))) + span.data = merged + + +def merge_refs_into_data(data: dict[str, Any] | None, refs: list[dict[str, Any]]) -> dict[str, Any]: + """Merge serialized refs into a span data dict, deduplicating with any present.""" + out = dict(data) if isinstance(data, dict) else {} + if refs: + existing = out.get(LINEAGE_REFS_KEY) + combined = list(existing) if isinstance(existing, list) else [] + combined.extend(refs) + out[LINEAGE_REFS_KEY] = _dedupe_dicts(combined) + return out + + +def _dedupe(refs: list[DataSourceRef]) -> list[dict[str, Any]]: + return _dedupe_dicts([ref.model_dump(exclude_none=True) for ref in refs]) + + +def _dedupe_dicts(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: + seen: set[tuple[Any, ...]] = set() + out: list[dict[str, Any]] = [] + for ref in refs: + key = (ref.get("namespace"), ref.get("name"), ref.get("version"), ref.get("role")) + if key not in seen: + seen.add(key) + out.append(ref) + return out diff --git a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index 6d186de5f..32b7bae73 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -65,6 +65,8 @@ def _add_source_to_span(span: Span, env_vars: EnvironmentVariables) -> None: span.data["__agent_name__"] = env_vars.AGENT_NAME if env_vars.AGENT_ID is not None: span.data["__agent_id__"] = env_vars.AGENT_ID + if env_vars.AGENT_VERSION is not None: + span.data["__agent_version__"] = env_vars.AGENT_VERSION def _build_sgp_span(span: Span, env_vars: EnvironmentVariables) -> SGPSpan: diff --git a/src/agentex/lib/environment_variables.py b/src/agentex/lib/environment_variables.py index 3113b78f4..cbad0f2d8 100644 --- a/src/agentex/lib/environment_variables.py +++ b/src/agentex/lib/environment_variables.py @@ -24,6 +24,7 @@ class EnvVarKeys(str, Enum): AGENT_NAME = "AGENT_NAME" AGENT_DESCRIPTION = "AGENT_DESCRIPTION" AGENT_ID = "AGENT_ID" + AGENT_VERSION = "AGENT_VERSION" AGENT_API_KEY = "AGENT_API_KEY" # ACP Configuration ACP_URL = "ACP_URL" @@ -66,6 +67,8 @@ class EnvironmentVariables(BaseModel): AGENT_NAME: str AGENT_DESCRIPTION: str | None = None AGENT_ID: str | None = None + # Build/version discriminator (image tag or git sha), set by the deployment + AGENT_VERSION: str | None = None AGENT_API_KEY: str | None = None ACP_TYPE: str | None = "async" AGENT_INPUT_TYPE: str | None = None diff --git a/src/agentex/resources/agents/schedules.py b/src/agentex/resources/agents/schedules.py index e6576b6d7..1750c513b 100644 --- a/src/agentex/resources/agents/schedules.py +++ b/src/agentex/resources/agents/schedules.py @@ -277,6 +277,7 @@ def list( self, agent_id: str, *, + include_live: bool | Omit = omit, limit: int | Omit = omit, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. @@ -289,6 +290,8 @@ def list( List run schedules for an agent. Args: + include_live: Include live Temporal state and upcoming action times. + extra_headers: Send extra headers extra_query: Add additional query parameters to the request @@ -306,7 +309,13 @@ def list( extra_query=extra_query, extra_body=extra_body, timeout=timeout, - query=maybe_transform({"limit": limit}, schedule_list_params.ScheduleListParams), + query=maybe_transform( + { + "include_live": include_live, + "limit": limit, + }, + schedule_list_params.ScheduleListParams, + ), ), cast_to=ScheduleListResponse, ) @@ -1056,6 +1065,7 @@ async def list( self, agent_id: str, *, + include_live: bool | Omit = omit, limit: int | Omit = omit, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. @@ -1068,6 +1078,8 @@ async def list( List run schedules for an agent. Args: + include_live: Include live Temporal state and upcoming action times. + extra_headers: Send extra headers extra_query: Add additional query parameters to the request @@ -1085,7 +1097,13 @@ async def list( extra_query=extra_query, extra_body=extra_body, timeout=timeout, - query=await async_maybe_transform({"limit": limit}, schedule_list_params.ScheduleListParams), + query=await async_maybe_transform( + { + "include_live": include_live, + "limit": limit, + }, + schedule_list_params.ScheduleListParams, + ), ), cast_to=ScheduleListResponse, ) diff --git a/src/agentex/types/agents/schedule_create_response.py b/src/agentex/types/agents/schedule_create_response.py index 4440629a2..9902bb832 100644 --- a/src/agentex/types/agents/schedule_create_response.py +++ b/src/agentex/types/agents/schedule_create_response.py @@ -86,6 +86,12 @@ class ScheduleCreateResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_list_params.py b/src/agentex/types/agents/schedule_list_params.py index 8a1d5561a..6e767e207 100644 --- a/src/agentex/types/agents/schedule_list_params.py +++ b/src/agentex/types/agents/schedule_list_params.py @@ -8,4 +8,7 @@ class ScheduleListParams(TypedDict, total=False): + include_live: bool + """Include live Temporal state and upcoming action times.""" + limit: int diff --git a/src/agentex/types/agents/schedule_list_response.py b/src/agentex/types/agents/schedule_list_response.py index 182bd53eb..faa19527f 100644 --- a/src/agentex/types/agents/schedule_list_response.py +++ b/src/agentex/types/agents/schedule_list_response.py @@ -86,6 +86,12 @@ class RunSchedule(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_pause_by_name_response.py b/src/agentex/types/agents/schedule_pause_by_name_response.py index 4c0ce1061..61232fdde 100644 --- a/src/agentex/types/agents/schedule_pause_by_name_response.py +++ b/src/agentex/types/agents/schedule_pause_by_name_response.py @@ -86,6 +86,12 @@ class SchedulePauseByNameResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_pause_response.py b/src/agentex/types/agents/schedule_pause_response.py index c14c83443..16e4a85ce 100644 --- a/src/agentex/types/agents/schedule_pause_response.py +++ b/src/agentex/types/agents/schedule_pause_response.py @@ -86,6 +86,12 @@ class SchedulePauseResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_resume_by_name_response.py b/src/agentex/types/agents/schedule_resume_by_name_response.py index 6bdbcd329..5999059c0 100644 --- a/src/agentex/types/agents/schedule_resume_by_name_response.py +++ b/src/agentex/types/agents/schedule_resume_by_name_response.py @@ -86,6 +86,12 @@ class ScheduleResumeByNameResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_resume_response.py b/src/agentex/types/agents/schedule_resume_response.py index 907792401..70e6e2aa0 100644 --- a/src/agentex/types/agents/schedule_resume_response.py +++ b/src/agentex/types/agents/schedule_resume_response.py @@ -86,6 +86,12 @@ class ScheduleResumeResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_retrieve_by_name_response.py b/src/agentex/types/agents/schedule_retrieve_by_name_response.py index 31663f41a..7b21ebb0a 100644 --- a/src/agentex/types/agents/schedule_retrieve_by_name_response.py +++ b/src/agentex/types/agents/schedule_retrieve_by_name_response.py @@ -86,6 +86,12 @@ class ScheduleRetrieveByNameResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_retrieve_response.py b/src/agentex/types/agents/schedule_retrieve_response.py index 20375347e..374ac635f 100644 --- a/src/agentex/types/agents/schedule_retrieve_response.py +++ b/src/agentex/types/agents/schedule_retrieve_response.py @@ -86,6 +86,12 @@ class ScheduleRetrieveResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_skip_response.py b/src/agentex/types/agents/schedule_skip_response.py index f4daf514b..f97216e0b 100644 --- a/src/agentex/types/agents/schedule_skip_response.py +++ b/src/agentex/types/agents/schedule_skip_response.py @@ -86,6 +86,12 @@ class ScheduleSkipResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_trigger_by_name_response.py b/src/agentex/types/agents/schedule_trigger_by_name_response.py index 036cf72f3..4006b7ad7 100644 --- a/src/agentex/types/agents/schedule_trigger_by_name_response.py +++ b/src/agentex/types/agents/schedule_trigger_by_name_response.py @@ -86,6 +86,12 @@ class ScheduleTriggerByNameResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_trigger_response.py b/src/agentex/types/agents/schedule_trigger_response.py index 22695f9c0..dbff651b0 100644 --- a/src/agentex/types/agents/schedule_trigger_response.py +++ b/src/agentex/types/agents/schedule_trigger_response.py @@ -86,6 +86,12 @@ class ScheduleTriggerResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_unskip_response.py b/src/agentex/types/agents/schedule_unskip_response.py index e25f31bbe..96f624220 100644 --- a/src/agentex/types/agents/schedule_unskip_response.py +++ b/src/agentex/types/agents/schedule_unskip_response.py @@ -86,6 +86,12 @@ class ScheduleUnskipResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_update_by_name_response.py b/src/agentex/types/agents/schedule_update_by_name_response.py index 8e8fd2112..e2905593c 100644 --- a/src/agentex/types/agents/schedule_update_by_name_response.py +++ b/src/agentex/types/agents/schedule_update_by_name_response.py @@ -86,6 +86,12 @@ class ScheduleUpdateByNameResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/src/agentex/types/agents/schedule_update_response.py b/src/agentex/types/agents/schedule_update_response.py index a27701d21..a7416e776 100644 --- a/src/agentex/types/agents/schedule_update_response.py +++ b/src/agentex/types/agents/schedule_update_response.py @@ -86,6 +86,12 @@ class ScheduleUpdateResponse(BaseModel): last_action_time: Optional[datetime] = None """When the schedule last fired.""" + live_data_available: Optional[bool] = None + """Whether requested live Temporal fields were retrieved successfully. + + Null when live enrichment was not requested. + """ + next_action_times: Optional[List[datetime]] = None """Upcoming scheduled fire times.""" diff --git a/tests/api_resources/agents/test_schedules.py b/tests/api_resources/agents/test_schedules.py index 362d32f77..8b281ee31 100644 --- a/tests/api_resources/agents/test_schedules.py +++ b/tests/api_resources/agents/test_schedules.py @@ -249,6 +249,7 @@ def test_method_list(self, client: Agentex) -> None: def test_method_list_with_all_params(self, client: Agentex) -> None: schedule = client.agents.schedules.list( agent_id="agent_id", + include_live=True, limit=1, ) assert_matches_type(ScheduleListResponse, schedule, path=["response"]) @@ -1205,6 +1206,7 @@ async def test_method_list(self, async_client: AsyncAgentex) -> None: async def test_method_list_with_all_params(self, async_client: AsyncAgentex) -> None: schedule = await async_client.agents.schedules.list( agent_id="agent_id", + include_live=True, limit=1, ) assert_matches_type(ScheduleListResponse, schedule, path=["response"]) diff --git a/tests/lib/core/harness/test_tracer_lineage.py b/tests/lib/core/harness/test_tracer_lineage.py new file mode 100644 index 000000000..75799caee --- /dev/null +++ b/tests/lib/core/harness/test_tracer_lineage.py @@ -0,0 +1,53 @@ +"""SpanTracer stamps registered data-source refs onto tool spans (SGP-6513).""" + +import pytest + +from agentex.lib.core.harness.types import OpenSpan, CloseSpan +from agentex.lib.core.harness.tracer import SpanTracer +from agentex.lib.core.tracing.lineage import ( + LINEAGE_REFS_KEY, + DataSourceRef, + clear_tool_sources, + register_tool_sources, +) + +from ._fakes import FakeTracing + + +@pytest.fixture(autouse=True) +def _clean_registry(): + clear_tool_sources() + yield + clear_tool_sources() + + +@pytest.mark.asyncio +async def test_tool_open_span_carries_registered_refs(): + register_tool_sources( + "query_guidance", + refs=[DataSourceRef("databricks://ey-tax", "guidance.rulings")], + resolver=lambda args: [DataSourceRef("elasticsearch://ey", args["index"])], + ) + fake = FakeTracing() + tracer = SpanTracer(trace_id="t1", parent_span_id="p1", tracing=fake) + + await tracer.handle(OpenSpan(key="c1", kind="tool", name="query_guidance", input={"index": "filings"})) + await tracer.handle(CloseSpan(key="c1", output={"ok": True}, is_complete=True)) + + (span,) = fake.ended_spans + namespaces = {ref["namespace"] for ref in span.data[LINEAGE_REFS_KEY]} + assert namespaces == {"databricks://ey-tax", "elasticsearch://ey"} + + +@pytest.mark.asyncio +async def test_unregistered_tool_and_reasoning_spans_carry_no_refs(): + fake = FakeTracing() + tracer = SpanTracer(trace_id="t1", parent_span_id=None, tracing=fake) + + await tracer.handle(OpenSpan(key="c1", kind="tool", name="unregistered", input={})) + await tracer.handle(CloseSpan(key="c1", output=None, is_complete=True)) + await tracer.handle(OpenSpan(key="reasoning:0", kind="reasoning", name="reasoning", input={})) + await tracer.handle(CloseSpan(key="reasoning:0", output="thought", is_complete=True)) + + for span in fake.ended_spans: + assert not (isinstance(span.data, dict) and LINEAGE_REFS_KEY in span.data) diff --git a/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py b/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py index dc8bab127..4a233fb72 100644 --- a/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py +++ b/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py @@ -39,6 +39,30 @@ def _make_mock_sgp_span() -> MagicMock: return sgp_span +class TestSourceStamps: + def test_agent_identity_and_version_stamped_into_span_data(self): + from agentex.lib.core.tracing.processors.sgp_tracing_processor import _add_source_to_span + + env = MagicMock(ACP_TYPE="async", AGENT_NAME="emu-tax", AGENT_ID="a1", AGENT_VERSION="sha-abc123") + span = _make_span() + _add_source_to_span(span, env) + assert span.data == { + "__source__": "agentex", + "__acp_type__": "async", + "__agent_name__": "emu-tax", + "__agent_id__": "a1", + "__agent_version__": "sha-abc123", + } + + def test_unset_identity_fields_are_omitted(self): + from agentex.lib.core.tracing.processors.sgp_tracing_processor import _add_source_to_span + + env = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) + span = _make_span() + _add_source_to_span(span, env) + assert span.data == {"__source__": "agentex"} + + # --------------------------------------------------------------------------- # Sync processor tests # --------------------------------------------------------------------------- @@ -48,7 +72,7 @@ class TestSGPSyncTracingProcessor: @staticmethod def _make_processor(): mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) mock_create_span = MagicMock(side_effect=lambda **kwargs: _make_mock_sgp_span()) with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.SGPClient"), patch( @@ -150,7 +174,7 @@ class TestSGPAsyncTracingProcessor: @staticmethod def _make_processor(): mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) mock_create_span = MagicMock(side_effect=lambda **kwargs: _make_mock_sgp_span()) mock_async_client = MagicMock() @@ -319,11 +343,9 @@ async def test_get_client_caches_per_event_loop(self): keepalive instead of paying a TLS handshake per span. """ mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) - with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch( - f"{MODULE}.AsyncSGPClient" - ) as mock_sgp_cls: + with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient") as mock_sgp_cls: mock_sgp_cls.side_effect = lambda **kwargs: MagicMock() from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( @@ -365,11 +387,11 @@ def capture_limits(*args, **kwargs): return original_async_client(*args, **kwargs) mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) - with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch( - f"{MODULE}.AsyncSGPClient" - ), patch("httpx.AsyncClient", side_effect=capture_limits): + with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient"), patch( + "httpx.AsyncClient", side_effect=capture_limits + ): from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( SGPAsyncTracingProcessor, ) @@ -380,8 +402,7 @@ def capture_limits(*args, **kwargs): assert len(captured_limits) == 1 max_keepalive = captured_limits[0].max_keepalive_connections assert max_keepalive is not None and max_keepalive > 0, ( - f"SGP async client should have keepalive enabled, got " - f"max_keepalive_connections={max_keepalive}" + f"SGP async client should have keepalive enabled, got max_keepalive_connections={max_keepalive}" ) def test_cache_is_weakkeydict_and_evicts_dead_loops(self): @@ -395,7 +416,7 @@ def test_cache_is_weakkeydict_and_evicts_dead_loops(self): import weakref mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient"): from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( @@ -428,18 +449,14 @@ async def test_disabled_processor_returns_none_client(self): from agentex.lib.types.tracing import SGPTracingProcessorConfig mock_env = MagicMock() - mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None) + mock_env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) - with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch( - f"{MODULE}.AsyncSGPClient" - ) as mock_sgp_cls: + with patch(f"{MODULE}.EnvironmentVariables", mock_env), patch(f"{MODULE}.AsyncSGPClient") as mock_sgp_cls: from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( SGPAsyncTracingProcessor, ) - processor = SGPAsyncTracingProcessor( - SGPTracingProcessorConfig(sgp_api_key="", sgp_account_id="") - ) + processor = SGPAsyncTracingProcessor(SGPTracingProcessorConfig(sgp_api_key="", sgp_account_id="")) assert processor._get_client() is None assert mock_sgp_cls.call_count == 0 diff --git a/tests/lib/core/tracing/test_lineage.py b/tests/lib/core/tracing/test_lineage.py new file mode 100644 index 000000000..c0fc3ebb9 --- /dev/null +++ b/tests/lib/core/tracing/test_lineage.py @@ -0,0 +1,147 @@ +"""Unit tests for the data-source ref module (sgp.lineage.refs capture).""" + +import json + +import pytest +from pydantic import ValidationError + +from agentex.lib.core.tracing.lineage import ( + LINEAGE_REFS_KEY, + DataSourceRef, + record, + data_sources, + resolve_refs, + clear_tool_sources, + merge_refs_into_data, + register_tool_sources, + resolve_refs_from_items, +) + + +@pytest.fixture(autouse=True) +def _clean_registry(): + clear_tool_sources() + yield + clear_tool_sources() + + +ES_REF = DataSourceRef("elasticsearch://ey-embryonic", "companies_v3") +DBX_REF = DataSourceRef("databricks://ey-tax", "guidance.rulings", role="input") + + +class TestDataSourceRef: + def test_positional_construction(self): + ref = DataSourceRef("s3://bucket", "key", version="v1", role="output") + assert ref.namespace == "s3://bucket" + assert ref.name == "key" + assert ref.version == "v1" + assert ref.role == "output" + + def test_non_uri_namespace_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("not-a-uri", "name") + + def test_underscore_host_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("mcp://ey_tax_server", "competitive-edge") + + def test_host_with_path_segment_allowed(self): + DataSourceRef("confluence://ey-tax/TAX", "page-123") + + def test_empty_name_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("s3://bucket", "") + + def test_bad_role_rejected(self): + with pytest.raises(ValidationError): + DataSourceRef("s3://bucket", "key", role="sideways") + + +class TestRegistryAndResolve: + def test_unregistered_tool_resolves_empty(self): + assert resolve_refs("unknown_tool", {}) == [] + + def test_static_refs(self): + register_tool_sources("search", refs=[ES_REF]) + refs = resolve_refs("search", {"q": "acme"}) + assert refs == [{"namespace": "elasticsearch://ey-embryonic", "name": "companies_v3", "role": "input"}] + + def test_resolver_refs_combined_with_static(self): + register_tool_sources( + "query_table", + refs=[ES_REF], + resolver=lambda args: [DataSourceRef("databricks://ey-tax", args["table"])], + ) + refs = resolve_refs("query_table", {"table": "guidance.rulings"}) + assert {r["namespace"] for r in refs} == {"elasticsearch://ey-embryonic", "databricks://ey-tax"} + + def test_resolver_failure_keeps_static_refs(self): + register_tool_sources("flaky", refs=[ES_REF], resolver=lambda args: args["missing"]) + refs = resolve_refs("flaky", {}) + assert len(refs) == 1 + + def test_reregistration_replaces(self): + register_tool_sources("search", refs=[ES_REF]) + register_tool_sources("search", refs=[DBX_REF]) + assert resolve_refs("search", {})[0]["namespace"] == "databricks://ey-tax" + + def test_dedupe(self): + register_tool_sources("search", refs=[ES_REF, ES_REF]) + assert len(resolve_refs("search", {})) == 1 + + +class TestDecorator: + def test_registers_by_function_name(self): + @data_sources(ES_REF) + def search_companies(q: str) -> str: + return q + + assert search_companies("x") == "x" + assert resolve_refs("search_companies", {}) != [] + + def test_registers_by_name_attribute(self): + class FakeFunctionTool: + name = "mcp_search" + + data_sources(DBX_REF)(FakeFunctionTool()) + assert resolve_refs("mcp_search", {}) != [] + + +class TestResolveFromItems: + def test_matches_function_call_items_and_parses_string_arguments(self): + register_tool_sources( + "query_table", + resolver=lambda args: [DataSourceRef("databricks://ey-tax", args["table"])], + ) + items = [ + {"type": "message", "content": []}, + {"type": "function_call", "name": "query_table", "arguments": json.dumps({"table": "t1"})}, + {"type": "function_call", "name": "unregistered", "arguments": "{}"}, + "not-a-dict", + ] + refs = resolve_refs_from_items(items) + assert refs == [{"namespace": "databricks://ey-tax", "name": "t1", "role": "input"}] + + def test_malformed_arguments_fall_back_to_static(self): + register_tool_sources("search", refs=[ES_REF]) + items = [{"type": "function_call", "name": "search", "arguments": "{not json"}] + assert len(resolve_refs_from_items(items)) == 1 + + +class TestRecordAndMerge: + def test_record_on_none_span_is_noop(self): + record(None, [ES_REF]) + + def test_record_merges_into_span_data(self): + class Span: + data = {"__span_type__": "CUSTOM"} + + span = Span() + record(span, [ES_REF]) + assert span.data["__span_type__"] == "CUSTOM" + assert span.data[LINEAGE_REFS_KEY][0]["name"] == "companies_v3" + + def test_merge_dedupes_against_existing(self): + data = merge_refs_into_data(None, [ES_REF.model_dump(exclude_none=True)]) + data = merge_refs_into_data(data, [ES_REF.model_dump(exclude_none=True)]) + assert len(data[LINEAGE_REFS_KEY]) == 1