diff --git a/.gitignore b/.gitignore index 3de6a9f..4e659c5 100644 --- a/.gitignore +++ b/.gitignore @@ -217,3 +217,7 @@ __marimo__/ # Pre-commit cache /.pre-commit-cache/ + +# Agents +# codex-instructions +local diff --git a/AGENTS.md b/AGENTS.md index 13e9fb5..40474f4 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -7,6 +7,7 @@ This document is a living map for developers building agents in the Codon ecosys - `docs/guides/codon-workload-quickstart.md` shows an end-to-end example of composing and running a workload directly with the SDK. - [`instrumentation-packages/codon-instrumentation-langgraph/AGENTS.md`](instrumentation-packages/codon-instrumentation-langgraph/AGENTS.md) covers the LangGraph decorators, the `LangGraphWorkloadAdapter`, and how to inherit telemetry automatically. - [`instrumentation-packages/codon-instrumentation-openai/AGENTS.md`](instrumentation-packages/codon-instrumentation-openai/AGENTS.md) will track OpenAI-specific instrumentation work. It currently outlines expectations and open tasks. +- Telemetry initialization now lives in the core SDK (`codon_sdk.instrumentation.initialize_telemetry`), with a hardcoded production default endpoint and API-key header support. Framework packages re-export this entrypoint to stay aligned. CodonWorkload can emit spans on its own when `enable_tracing=True` (default: False); leave it disabled if another instrumentation layer is already wrapping nodes to avoid duplicate spans. If you are introducing a new framework integration, create an `AGENTS.md` alongside it and link back here. @@ -21,6 +22,7 @@ Telemetry defaults to OpenTelemetry OTLP exporters; align your environment varia - The OpenAI instrumentation package is still a stub—see its `AGENTS.md` for the punch list. - SDK testing and documentation are light; expect breaking changes as the contracts tighten. - Additional frameworks (e.g., LangChain, LiteLLM) may join this layout. Document them here when they land. +- Telemetry bootstrap is centralized; callers can override endpoint/API key via args or env vars (`OTEL_EXPORTER_OTLP_ENDPOINT`, `CODON_API_KEY`). An opt-in flag (`CODON_ATTACH_TO_EXISTING_OTEL_PROVIDER`) lets Codon attach its OTLP exporter to an existing tracer provider (e.g., when users already run OTEL auto-instrumentation) instead of replacing it. ## Contribution Guidelines - Keep the `AGENTS.md` files in sync with code changes that alter agent-facing behavior. diff --git a/codex-instructions/sdk-production-refactor.md b/codex-instructions/sdk-production-refactor.md new file mode 100644 index 0000000..b5256ba --- /dev/null +++ b/codex-instructions/sdk-production-refactor.md @@ -0,0 +1,174 @@ +### Phase 1: The "Front Door" (API Gateway) + +**Objective:** Secure the ingestion pipeline immediately with minimal complexity. +**Timeline:** Immediate / MVP. + +#### 1\. SDK Changes (Core Refactoring) + + * **Move Logic:** Move `initialize_telemetry` from `codon.instrumentation.langgraph` to `codon_sdk.instrumentation`. + * **Update Signature:** Accept `api_key` and `endpoint` arguments. + * **Hardcode Default:** Set the default endpoint to your production load balancer (e.g., `https://ingest.codonops.ai:4317`). + * **Header Injection:** Inject `x-codon-api-key` into the exporter. + +**Action:** Create `sdk/src/codon_sdk/instrumentation/config.py` (or similar) and migrate the logic there. Update the LangGraph package to import it. + +#### 2\. Infrastructure Changes (Server-Side) + + * **DNS:** Point `ingest.codonops.ai` to your Cloud Load Balancer. + * **Gateway:** Deploy Nginx (or Kong) behind the LB. + * **Config:** specific rule to check `x-codon-api-key` against your DB/Cache. + * **Success:** If valid, inject header `X-Codon-Org-ID: ` and forward to Collector. + * **Failure:** Return `401 Unauthorized`. + * **OTel Collector:** + * Enable `include_metadata: true` on OTLP receivers. + * Add `attributes/insert_tenancy` processor to extract `metadata.x-codon-org-id`. + +----- + +### Phase 2: The "Infinite Scale" (JWT) + +**Objective:** Decouple ingestion authorization from the database to handle high-volume telemetry. +**Timeline:** Post-Launch / Scale-up. + +#### 1\. SDK Changes + + * **Token Manager:** Implement a `TokenManager` class in the Core SDK. + * **Exchange:** On init, call `https://api.codon.ai/auth/token` with the API Key to get a JWT. + * **Refresh:** Automatically refresh the JWT 5 minutes before expiry. + * **Exporter Update:** Update `initialize_telemetry` to use the dynamic JWT in the `Authorization: Bearer ` header instead of the static API key. + +#### 2\. Infrastructure Changes + + * **Auth Service:** Ensure your backend has an endpoint to mint signed JWTs (RS256) containing the `org_id` claim. + * **Collector/Gateway:** Reconfigure to use the **OIDC Authenticator**. + * It will now validate the **signature** of the JWT locally (CPU operation) instead of calling the DB (I/O operation). + * The `attributes` processor will now extract tenancy from `auth.claims.org_id`. + +----- + +### Architecture Decision Records (ADRs) + +Here are the formal records you should commit to your documentation (e.g., `docs/adr/`). + +#### ADR 001: Telemetry Authentication Strategy + +**Status:** Accepted +**Date:** 2025-11-24 + +**Context:** +We need to secure the OpenTelemetry ingestion endpoint (`ingest.codonops.ai`) to prevent unauthorized data submission and ensure accurate billing/tenancy. We have two options: checking API keys against a database (Gateway pattern) or using signed tokens (JWT pattern). + +**Decision:** +We will implement a **Two-Phase Strategy**: + +1. **Phase 1 (API Gateway):** Validate API Keys at the ingress gateway (Nginx) via a database/cache lookup. +2. **Phase 2 (JWT):** Migrate to client-side JWT exchange for stateless authentication at scale. + +**Consequences:** + + * **Positive:** Phase 1 allows for rapid implementation using existing API keys without complex SDK logic. Phase 2 provides a clear path to handling massive concurrency without database bottlenecks. + * **Negative:** Phase 1 introduces a database dependency on the hot path of telemetry ingestion. If the DB slows down, telemetry ingestion slows down. This is acceptable for MVP volumes but must be replaced (Phase 2) before high-scale adoption. + +#### ADR 002: Centralization of Telemetry Configuration + +**Status:** Accepted +**Date:** 2025-11-24 + +**Context:** +Currently, `initialize_telemetry` is defined inside the `codon-instrumentation-langgraph` package. As we add support for OpenAI, CrewAI, and other frameworks, duplicating this initialization logic will lead to drift in configuration defaults (endpoints, auth headers, resource attributes). + +**Decision:** +We will move the telemetry initialization logic to the **SDK Core** (`codon_sdk`). + + * **New Location:** `codon_sdk.instrumentation.initialize` (or similar). + * **Instrumentation Packages:** Will import and wrap this core function or instruct users to call it directly. + +**Consequences:** + + * **Positive:** Single source of truth for endpoints and authentication. Updates to the ingestion architecture (e.g., changing the endpoint URL or auth header format) only need to happen in one place (`codon_sdk`). + * **Negative:** Instrumentation packages take a hard dependency on the specific version of `codon_sdk` that contains this utility. + +----- + +### Recommended Code Structure for Core SDK + +Here is how you should implement the **Core SDK** change right now to support Phase 1. + +**File:** `sdk/src/codon_sdk/instrumentation/__init__.py` (or a new `config.py` inside it) + +```python +import os +from typing import Optional, Dict +from opentelemetry import trace +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter + +# Hardcoded Production Endpoint - The "Front Door" +DEFAULT_INGEST_ENDPOINT = "https://ingest.codonops.ai:4317" + +def initialize_telemetry( + api_key: Optional[str] = None, + service_name: Optional[str] = None, + endpoint: Optional[str] = None, +) -> None: + """ + Global initialization for Codon telemetry. + + This should be called once at the start of the application. It configures + OpenTelemetry to send traces to the Codon Cloud (or a custom collector). + """ + + # 1. Resolve Identity + final_api_key = api_key or os.getenv("CODON_API_KEY") + if not final_api_key: + # Fallback for local dev without auth (optional, or raise Error) + # For Phase 1 production, we might want to enforce this. + pass + + # 2. Resolve Context + # We prefer the user-provided service name, then env var, then default. + final_service_name = ( + service_name + or os.getenv("OTEL_SERVICE_NAME") + or "unknown_codon_service" + ) + + # 3. Resolve Destination + # Default to Prod, allow env var override, allow explicit arg override. + final_endpoint = ( + endpoint + or os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT") + or DEFAULT_INGEST_ENDPOINT + ) + + # 4. Configure Headers + headers: Dict[str, str] = {} + if final_api_key: + headers["x-codon-api-key"] = final_api_key + + # 5. Setup OpenTelemetry + resource = Resource(attributes={"service.name": final_service_name}) + + exporter = OTLPSpanExporter( + endpoint=final_endpoint, + headers=headers + ) + + provider = TracerProvider(resource=resource) + provider.add_span_processor(BatchSpanProcessor(exporter)) + + # Set the global tracer provider so all instrumentation packages pick it up + trace.set_tracer_provider(provider) +``` + +**Usage in `codon-instrumentation-langgraph`:** +You would delete the local `initialize_telemetry` definition and re-export or document the usage of the core one. + +```python +# instrumentation-packages/codon-instrumentation-langgraph/codon/instrumentation/langgraph/__init__.py + +# Re-export for convenience, or deprecate and point users to core +from codon_sdk.instrumentation import initialize_telemetry +``` diff --git a/instrumentation-packages/codon-instrumentation-langgraph/AGENTS.md b/instrumentation-packages/codon-instrumentation-langgraph/AGENTS.md index 4cc07a1..a312b97 100644 --- a/instrumentation-packages/codon-instrumentation-langgraph/AGENTS.md +++ b/instrumentation-packages/codon-instrumentation-langgraph/AGENTS.md @@ -160,7 +160,7 @@ The ledger records each iteration through the loop, and `runtime.state` tracks i --- ## Telemetry & Audit Integration -- Call `initialize_telemetry(service_name=...)` once during process startup to export spans via OTLP. +- Call `initialize_telemetry(service_name=...)` once during process startup to export spans via OTLP. The initializer now lives in the core SDK (`codon_sdk.instrumentation.initialize_telemetry`) and is re-exported here. It defaults the endpoint to `https://ingest.codonops.ai:4317`, injects `x-codon-api-key` from the argument or `CODON_API_KEY` env, and respects `OTEL_EXPORTER_OTLP_ENDPOINT`/`OTEL_SERVICE_NAME` overrides. If you already have an OTEL tracer provider (e.g., via auto-instrumentation), set `CODON_ATTACH_TO_EXISTING_OTEL_PROVIDER=true` or pass `attach_to_existing=True` to add Codon’s exporter to the existing provider instead of replacing it. - Each node span now carries workload metadata (`codon.workload.id`, `codon.workload.run_id`, `codon.workload.logic_id`, `codon.workload.deployment_id`, `codon.organization.id`) so traces can be rolled up by workload, deployment, or organization without joins. - `LangGraphTelemetryCallback` is attached automatically when invoking LangChain runnables; it captures model vendor/identifier, token usage (prompt, completion, total), and response metadata, all of which is emitted as span attributes (`codon.tokens.*`, `codon.model.*`, `codon.node.raw_attributes_json`). - Instrumentation writes into the shared `NodeTelemetryPayload` (`runtime.telemetry`) defined by the SDK so future mixins collect the same schema-aligned fields without reimplementing bookkeeping. diff --git a/instrumentation-packages/codon-instrumentation-langgraph/codon/instrumentation/langgraph/__init__.py b/instrumentation-packages/codon-instrumentation-langgraph/codon/instrumentation/langgraph/__init__.py index 183ac2f..bee11fe 100644 --- a/instrumentation-packages/codon-instrumentation-langgraph/codon/instrumentation/langgraph/__init__.py +++ b/instrumentation-packages/codon-instrumentation-langgraph/codon/instrumentation/langgraph/__init__.py @@ -23,16 +23,13 @@ from opentelemetry import trace from opentelemetry.trace import Status, StatusCode -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import BatchSpanProcessor -from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter -from opentelemetry.sdk.resources import Resource from .attributes import LangGraphSpanAttributes from codon_sdk.instrumentation.schemas.nodespec import NodeSpec, NodeSpecSpanAttributes from codon_sdk.agents import Workload from codon_sdk.instrumentation.schemas.telemetry.spans import CodonBaseSpanAttributes from codon_sdk.instrumentation.telemetry import NodeTelemetryPayload +from codon_sdk.instrumentation import initialize_telemetry __all__ = [ "LangGraphWorkloadMixin", @@ -45,7 +42,6 @@ "LangGraphTelemetryCallback", ] -SERVICE_NAME: str = os.getenv("OTEL_SERVICE_NAME") ORG_NAMESPACE: str = os.getenv("ORG_NAMESPACE") __framework__ = "langgraph" @@ -78,31 +74,6 @@ def from_langgraph( """Translate a LangGraph graph into a concrete Codon workload.""" -def initialize_telemetry(service_name: str = SERVICE_NAME or "") -> None: - """Initialize OpenTelemetry tracing with OTLP exporter. - - Sets up a TracerProvider with BatchSpanProcessor and OTLPSpanExporter to - export spans to an OTLP collector endpoint. - - Args: - service_name: Name for the service in telemetry data. Defaults to - OTEL_SERVICE_NAME environment variable or empty string. - - Example: - >>> initialize_telemetry(service_name="codon-langgraph-demo") - """ - # Set up service name - resource = Resource(attributes={"service.name": service_name}) - # Set up TracerProvider & Exporter - provider = TracerProvider(resource=resource) - otlp_exporter = OTLPSpanExporter() - - processor = BatchSpanProcessor(otlp_exporter) - provider.add_span_processor(processor) - - trace.set_tracer_provider(provider) - - def current_invocation() -> Optional[NodeTelemetryPayload]: """Return the currently active node invocation telemetry (if any).""" @@ -352,7 +323,7 @@ def track_node( When the decorated function executes, this decorator: - Materializes a NodeSpec and captures its ID, signature, and schemas - - Wraps execution in an OpenTelemetry span (async and sync supported) + - Wraps execution in an OpenTelemetry span (async and sync supported) - Records inputs, outputs, and wall-clock latency via standardized span attributes Args: diff --git a/sdk/AGENTS.md b/sdk/AGENTS.md index 5533520..22011a5 100644 --- a/sdk/AGENTS.md +++ b/sdk/AGENTS.md @@ -82,6 +82,8 @@ For a full walkthrough, see `docs/guides/codon-workload-quickstart.md`. ## Telemetry Vocabulary - `codon_sdk.instrumentation.schemas.telemetry.spans` defines span names and base attributes shared across frameworks. - Document additions here before using them in instrumentation packages to keep alignment. +- Telemetry initialization is centralized in `codon_sdk.instrumentation.initialize_telemetry`, with default endpoint `https://ingest.codonops.ai:4317` and `x-codon-api-key` header support (args override env; env vars `OTEL_EXPORTER_OTLP_ENDPOINT`, `CODON_API_KEY`, `OTEL_SERVICE_NAME` remain valid). Optional attach mode (`attach_to_existing` arg or `CODON_ATTACH_TO_EXISTING_OTEL_PROVIDER` env) lets you add Codon’s exporter to an existing tracer provider instead of replacing it—useful when OTEL auto-instrumentation is already active. +- CodonWorkload can emit spans natively when `enable_tracing=True` (default: False). It uses the global tracer provider configured via `initialize_telemetry` to create one span per node execution with workload/org/deployment IDs, logic/run IDs, and NodeSpec attributes. Leave it disabled if another instrumentation layer (e.g., LangGraph adapter) is already wrapping nodes to avoid duplicate spans. ## Extending the SDK - Capture requirements for new schema fields inside each class docstring and mirror them here. diff --git a/sdk/pyproject.toml b/sdk/pyproject.toml index 40b37a7..4d10c95 100644 --- a/sdk/pyproject.toml +++ b/sdk/pyproject.toml @@ -27,6 +27,9 @@ classifiers = [ ] # Core dependencies needed by every part of the SDK dependencies = [ + "opentelemetry-api", + "opentelemetry-exporter-otlp-proto-grpc", + "opentelemetry-sdk", "pydantic", ] diff --git a/sdk/src/codon_sdk/agents/codon_workload.py b/sdk/src/codon_sdk/agents/codon_workload.py index 3c65f27..e376809 100644 --- a/sdk/src/codon_sdk/agents/codon_workload.py +++ b/sdk/src/codon_sdk/agents/codon_workload.py @@ -42,6 +42,9 @@ ) from uuid import uuid4 +from opentelemetry import trace +from opentelemetry.trace import Status, StatusCode + from codon_sdk.instrumentation.schemas.logic_id import ( AgentClass, LogicRequest, @@ -49,8 +52,9 @@ Topology, generate_logic_id, ) -from codon_sdk.instrumentation.schemas.nodespec import NodeSpec +from codon_sdk.instrumentation.schemas.nodespec import NodeSpec, NodeSpecSpanAttributes from codon_sdk.instrumentation.telemetry import NodeTelemetryPayload +from codon_sdk.instrumentation.schemas.telemetry.spans import CodonBaseSpanAttributes from .workload import Workload @@ -68,6 +72,98 @@ def _render_payload(value: Any, *, max_length: int = 2048) -> str: return rendered +def _apply_nodespec_attributes(span, nodespec: NodeSpec) -> None: + span.set_attribute(NodeSpecSpanAttributes.ID.value, nodespec.id) + span.set_attribute(NodeSpecSpanAttributes.Version.value, nodespec.spec_version) + span.set_attribute(NodeSpecSpanAttributes.Name.value, nodespec.name) + span.set_attribute(NodeSpecSpanAttributes.Role.value, nodespec.role) + span.set_attribute( + NodeSpecSpanAttributes.CallableSignature.value, + nodespec.callable_signature, + ) + span.set_attribute( + NodeSpecSpanAttributes.InputSchema.value, + nodespec.input_schema, + ) + if nodespec.output_schema is not None: + span.set_attribute( + NodeSpecSpanAttributes.OutputSchema.value, + nodespec.output_schema, + ) + if nodespec.model_name: + span.set_attribute(NodeSpecSpanAttributes.ModelName.value, nodespec.model_name) + if nodespec.model_version: + span.set_attribute( + NodeSpecSpanAttributes.ModelVersion.value, + nodespec.model_version, + ) + + +def _apply_workload_attributes( + span, *, telemetry: NodeTelemetryPayload, nodespec: NodeSpec, context: Dict[str, Any] +) -> None: + span.set_attribute(CodonBaseSpanAttributes.AgentFramework.value, "codon") + organization = ( + telemetry.organization_id + or telemetry.org_namespace + or context.get("organization_id") + or context.get("org_namespace") + or nodespec.org_namespace + or os.getenv("ORG_NAMESPACE") + ) + if organization: + span.set_attribute(CodonBaseSpanAttributes.OrganizationId.value, organization) + span.set_attribute(CodonBaseSpanAttributes.OrgNamespace.value, organization) + + workload_id = telemetry.workload_id or context.get("workload_id") + logic_id = telemetry.workload_logic_id or context.get("logic_id") + run_id = telemetry.workload_run_id or context.get("workload_run_id") + deployment_id = telemetry.deployment_id or context.get("deployment_id") + + if workload_id: + span.set_attribute(CodonBaseSpanAttributes.WorkloadId.value, workload_id) + if logic_id: + span.set_attribute(CodonBaseSpanAttributes.WorkloadLogicId.value, logic_id) + if run_id: + span.set_attribute(CodonBaseSpanAttributes.WorkloadRunId.value, run_id) + if deployment_id: + span.set_attribute(CodonBaseSpanAttributes.DeploymentId.value, deployment_id) + + if telemetry.workload_name or context.get("workload_name"): + span.set_attribute( + CodonBaseSpanAttributes.WorkloadName.value, + telemetry.workload_name or context.get("workload_name"), + ) + if telemetry.workload_version or context.get("workload_version"): + span.set_attribute( + CodonBaseSpanAttributes.WorkloadVersion.value, + telemetry.workload_version or context.get("workload_version"), + ) + + span.set_attribute(CodonBaseSpanAttributes.NodeStatusCode.value, telemetry.status_code) + if telemetry.error_message: + span.set_attribute(CodonBaseSpanAttributes.NodeErrorMessage.value, telemetry.error_message) + + if telemetry.node_input is not None: + span.set_attribute(CodonBaseSpanAttributes.NodeInput.value, telemetry.node_input) + if telemetry.node_output is not None: + span.set_attribute(CodonBaseSpanAttributes.NodeOutput.value, telemetry.node_output) + if telemetry.duration_ms is not None: + span.set_attribute(CodonBaseSpanAttributes.NodeLatencyMs.value, telemetry.duration_ms) + + if telemetry.model_vendor: + span.set_attribute(CodonBaseSpanAttributes.ModelVendor.value, telemetry.model_vendor) + if telemetry.model_identifier: + span.set_attribute(CodonBaseSpanAttributes.ModelIdentifier.value, telemetry.model_identifier) + + if telemetry.input_tokens is not None: + span.set_attribute(CodonBaseSpanAttributes.TokenInput.value, telemetry.input_tokens) + if telemetry.output_tokens is not None: + span.set_attribute(CodonBaseSpanAttributes.TokenOutput.value, telemetry.output_tokens) + if telemetry.total_tokens is not None: + span.set_attribute(CodonBaseSpanAttributes.TokenTotal.value, telemetry.total_tokens) + + @dataclass(frozen=True) class Token: """A unit of work travelling between nodes.""" @@ -282,6 +378,7 @@ def __init__( version: str, description: Optional[str] = None, tags: Optional[Sequence[str]] = None, + enable_tracing: bool = False, ) -> None: """Initialize workload with metadata. @@ -290,6 +387,9 @@ def __init__( version: Semantic version (e.g., "1.2.0"). description: Human-readable purpose description. tags: Categorization tags. + enable_tracing: When True, emit OpenTelemetry spans directly from the + CodonWorkload runtime. Defaults to False to avoid double + instrumentation when external wrappers (e.g., LangGraph) are used. """ self._node_specs: Dict[str, NodeSpec] = {} self._node_functions: Dict[str, Callable[..., Any]] = {} @@ -300,6 +400,7 @@ def __init__( self._logic_id: Optional[str] = None self._entry_nodes: Optional[List[str]] = None self._organization_id: Optional[str] = os.getenv("ORG_NAMESPACE") + self._enable_tracing = enable_tracing super().__init__( name=name, version=version, @@ -393,7 +494,7 @@ def add_edge(self, source_name: str, destination_name: str) -> None: Raises: WorkloadRegistrationError: If either node name doesn't exist. - + TODO: Clarify what source_name and destination_name semantically represent TODO: Document whether this creates directed edges and how token flow works """ @@ -462,6 +563,7 @@ async def execute_async( state: Dict[str, Any] = {} loop = asyncio.get_running_loop() + tracer = trace.get_tracer(__name__) if self._enable_tracing else None async def publish_event(event_type: str, **data: Any) -> None: if event_handler is None: @@ -573,44 +675,59 @@ def enqueue( ) started_at = _utcnow() + span_ctx = ( + tracer.start_as_current_span(f"codon.node.{node_name}") + if tracer + else contextlib.nullcontext() + ) await publish_event( "node_started", node=node_name, token_id=token.id, payload=token.payload, ) - try: - result = node_callable(token.payload, runtime=handle, context=context) - if inspect.isawaitable(result): - result = await result # type: ignore[assignment] - except Exception as exc: # pragma: no cover - telemetry.status_code = "ERROR" - telemetry.error_message = repr(exc) - ledger.append( - AuditEvent( - event_type="node_failed", - timestamp=_utcnow(), - token_id=token.id, - source_node=node_name, - target_node=None, - metadata={"exception": repr(exc)}, + with span_ctx as span: + if span is not None: + _apply_nodespec_attributes(span, nodespec) + _apply_workload_attributes(span, telemetry=telemetry, nodespec=nodespec, context=context) + try: + result = node_callable(token.payload, runtime=handle, context=context) + if inspect.isawaitable(result): + result = await result # type: ignore[assignment] + except Exception as exc: # pragma: no cover + telemetry.status_code = "ERROR" + telemetry.error_message = repr(exc) + if span is not None: + _apply_workload_attributes(span, telemetry=telemetry, nodespec=nodespec, context=context) + span.set_status(Status(StatusCode.ERROR, str(exc))) + ledger.append( + AuditEvent( + event_type="node_failed", + timestamp=_utcnow(), + token_id=token.id, + source_node=node_name, + target_node=None, + metadata={"exception": repr(exc)}, + ) ) - ) - schedule_event( - "node_failed", - { - "node": node_name, - "token_id": token.id, - "exception": repr(exc), - }, - ) - raise WorkloadRuntimeError( - f"Node '{node_name}' execution failed" - ) from exc - finished_at = _utcnow() - telemetry.duration_ms = int((finished_at - started_at).total_seconds() * 1000) - telemetry.status_code = "OK" - telemetry.node_output = _render_payload(result) + schedule_event( + "node_failed", + { + "node": node_name, + "token_id": token.id, + "exception": repr(exc), + }, + ) + raise WorkloadRuntimeError( + f"Node '{node_name}' execution failed" + ) from exc + finished_at = _utcnow() + telemetry.duration_ms = int((finished_at - started_at).total_seconds() * 1000) + telemetry.status_code = "OK" + telemetry.node_output = _render_payload(result) + if span is not None: + _apply_workload_attributes(span, telemetry=telemetry, nodespec=nodespec, context=context) + span.set_status(Status(StatusCode.OK)) records[node_name].append( NodeExecutionRecord( diff --git a/sdk/src/codon_sdk/instrumentation/__init__.py b/sdk/src/codon_sdk/instrumentation/__init__.py index fc237f7..6bdab92 100644 --- a/sdk/src/codon_sdk/instrumentation/__init__.py +++ b/sdk/src/codon_sdk/instrumentation/__init__.py @@ -11,3 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. + +from .config import DEFAULT_INGEST_ENDPOINT, initialize_telemetry + +__all__ = ["DEFAULT_INGEST_ENDPOINT", "initialize_telemetry"] diff --git a/sdk/src/codon_sdk/instrumentation/config.py b/sdk/src/codon_sdk/instrumentation/config.py new file mode 100644 index 0000000..575121e --- /dev/null +++ b/sdk/src/codon_sdk/instrumentation/config.py @@ -0,0 +1,113 @@ +# Copyright 2025 Codon, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from typing import Dict, Optional + +from opentelemetry import trace +from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor + +import logging + +# Avoid configuring root logger; module-level logger only. +logger = logging.getLogger(__name__) + +# Hardcoded production ingest endpoint; can be overridden via argument or env. +DEFAULT_INGEST_ENDPOINT = "https://ingest.codonops.ai:4317" + + +_TRUE_VALUES = {"1", "true", "t", "yes", "y", "on"} + + +def _coerce_bool(value: Optional[str]) -> Optional[bool]: + if value is None: + return None + return value.strip().lower() in _TRUE_VALUES + + +def initialize_telemetry( + api_key: Optional[str] = None, + service_name: Optional[str] = None, + endpoint: Optional[str] = None, + attach_to_existing: Optional[bool] = None, +) -> None: + """Initialize OpenTelemetry tracing for Codon. + + Endpoint precedence: explicit argument, ``OTEL_EXPORTER_OTLP_ENDPOINT``, then + production default. API key precedence: explicit argument, then + ``CODON_API_KEY`` environment variable. When provided, the API key is sent + as ``x-codon-api-key`` on OTLP requests. + """ + + attach = ( + attach_to_existing + if attach_to_existing is not None + else _coerce_bool(os.getenv("CODON_ATTACH_TO_EXISTING_OTEL_PROVIDER")) + ) or False + + existing_provider = trace.get_tracer_provider() + + final_api_key = api_key or os.getenv("CODON_API_KEY") + final_service_name = ( + service_name + or os.getenv("OTEL_SERVICE_NAME") + or "unknown_codon_service" + ) + final_endpoint = ( + endpoint + or os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT") + or DEFAULT_INGEST_ENDPOINT + ) + + headers: Dict[str, str] = {} + if final_api_key: + headers["x-codon-api-key"] = final_api_key + else: + logger.warning( + "CODON telemetry initialized without an API key; spans may be rejected by the gateway" + ) + + resource = Resource(attributes={"service.name": final_service_name}) + exporter = OTLPSpanExporter(endpoint=final_endpoint, headers=headers) + + if attach and isinstance(existing_provider, TracerProvider): + processor = BatchSpanProcessor(exporter) + # Avoid double-adding an equivalent processor if initialise is called repeatedly. + if not _has_equivalent_processor(existing_provider, exporter): + existing_provider.add_span_processor(processor) + return + + provider = TracerProvider(resource=resource) + provider.add_span_processor(BatchSpanProcessor(exporter)) + + trace.set_tracer_provider(provider) + + +def _has_equivalent_processor(provider: TracerProvider, exporter: OTLPSpanExporter) -> bool: + """Return True if provider already has a processor targeting the same exporter endpoint/headers.""" + processors = getattr(getattr(provider, "_active_span_processor", None), "processors", None) + if not processors: + return False + for processor in processors: + existing_exporter = getattr(processor, "span_exporter", None) + if existing_exporter is exporter: + return True + same_endpoint = getattr(existing_exporter, "endpoint", None) == getattr(exporter, "endpoint", None) + same_headers = getattr(existing_exporter, "headers", None) == getattr(exporter, "headers", None) + if same_endpoint and same_headers: + return True + return False diff --git a/sdk/test/agents/test_codon_workload.py b/sdk/test/agents/test_codon_workload.py index a57075d..e46ef7a 100644 --- a/sdk/test/agents/test_codon_workload.py +++ b/sdk/test/agents/test_codon_workload.py @@ -198,6 +198,78 @@ def node(message, *, runtime, context): assert captured_context["org_namespace"] == "context-org" assert captured_context["workload_name"] == "ContextAgent" assert captured_context["workload_version"] == "1.0.0" + + +def test_tracing_opt_in_emits_spans(monkeypatch): + opentelemetry = pytest.importorskip("opentelemetry") + sdk_trace = pytest.importorskip("opentelemetry.sdk.trace") + export = pytest.importorskip("opentelemetry.sdk.trace.export") + + InMemorySpanExporter = export.InMemorySpanExporter + SimpleSpanProcessor = export.SimpleSpanProcessor + TracerProvider = sdk_trace.TracerProvider + trace = opentelemetry.trace + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + trace.set_tracer_provider(provider) + + workload = CodonWorkload(name="TraceAgent", version="1.0.0", enable_tracing=True) + + def first(message, *, runtime, context): + runtime.emit("second", {"value": message["value"] + 1}) + return {"value": message["value"]} + + def second(message, *, runtime, context): + return message["value"] + + workload.add_node(first, name="first", role="starter") + workload.add_node(second, name="second", role="finisher") + workload.add_edge("first", "second") + + workload.execute({"value": 1}, deployment_id="dev-trace") + + spans = exporter.get_finished_spans() + assert len(spans) == 2 + span_names = {span.name for span in spans} + assert "codon.node.first" in span_names + assert "codon.node.second" in span_names + for span in spans: + assert span.attributes.get("codon.workload.id") == workload.agent_class_id + assert span.attributes.get("codon.workload.logic_id") == workload.logic_id + assert span.attributes.get("codon.workload.run_id") + assert span.attributes.get("codon.workload.deployment_id") == "dev-trace" + exporter.clear() + + +def test_tracing_opt_out_produces_no_spans(monkeypatch): + opentelemetry = pytest.importorskip("opentelemetry") + sdk_trace = pytest.importorskip("opentelemetry.sdk.trace") + export = pytest.importorskip("opentelemetry.sdk.trace.export") + + InMemorySpanExporter = export.InMemorySpanExporter + SimpleSpanProcessor = export.SimpleSpanProcessor + TracerProvider = sdk_trace.TracerProvider + trace = opentelemetry.trace + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + trace.set_tracer_provider(provider) + + workload = CodonWorkload(name="TraceOff", version="1.0.0", enable_tracing=False) + + def node(message, *, runtime, context): + return message["value"] + + workload.add_node(node, name="only", role="solo") + + workload.execute({"value": 1}, deployment_id="dev") + + spans = exporter.get_finished_spans() + assert len(spans) == 0 + exporter.clear() assert captured_context["deployment_id"] == "dev-west" diff --git a/sdk/test/instrumentation/langgraph/test_initialize_reexport.py b/sdk/test/instrumentation/langgraph/test_initialize_reexport.py new file mode 100644 index 0000000..65c213f --- /dev/null +++ b/sdk/test/instrumentation/langgraph/test_initialize_reexport.py @@ -0,0 +1,8 @@ +import pytest + +core_module = pytest.importorskip("codon_sdk.instrumentation.config") +langgraph_module = pytest.importorskip("codon.instrumentation.langgraph") + + +def test_initialize_telemetry_reexported(): + assert langgraph_module.initialize_telemetry is core_module.initialize_telemetry diff --git a/sdk/test/instrumentation/test_initialize_telemetry.py b/sdk/test/instrumentation/test_initialize_telemetry.py new file mode 100644 index 0000000..296f6cf --- /dev/null +++ b/sdk/test/instrumentation/test_initialize_telemetry.py @@ -0,0 +1,178 @@ +import pytest +from types import SimpleNamespace +import logging + +from codon_sdk.instrumentation import DEFAULT_INGEST_ENDPOINT +from codon_sdk.instrumentation import config as instrumentation_config + + +class DummyResource: + def __init__(self, *, attributes): + self.attributes = attributes + + +class DummyExporter: + def __init__(self, *, endpoint=None, headers=None, **kwargs): + self.endpoint = endpoint + self.headers = headers or {} + + +class DummyProcessor: + def __init__(self, exporter): + self.span_exporter = exporter + + +class DummyProvider: + def __init__(self, resource=None): + self.resource = resource or DummyResource(attributes={}) + self.processors = [] + self._active_span_processor = SimpleNamespace(processors=self.processors) + + def add_span_processor(self, processor): + self.processors.append(processor) + + +@pytest.fixture(autouse=True) +def clear_env(monkeypatch): + for key in [ + "CODON_API_KEY", + "OTEL_EXPORTER_OTLP_ENDPOINT", + "OTEL_SERVICE_NAME", + "CODON_ATTACH_TO_EXISTING_OTEL_PROVIDER", + ]: + monkeypatch.delenv(key, raising=False) + yield + + +def _patch_base(monkeypatch, existing_provider=None): + monkeypatch.setattr(instrumentation_config, "OTLPSpanExporter", DummyExporter) + monkeypatch.setattr(instrumentation_config, "BatchSpanProcessor", DummyProcessor) + monkeypatch.setattr(instrumentation_config, "TracerProvider", DummyProvider) + monkeypatch.setattr(instrumentation_config, "Resource", DummyResource) + monkeypatch.setattr( + instrumentation_config.trace, + "get_tracer_provider", + lambda: existing_provider, + ) + # Ensure warnings are captured + monkeypatch.setattr(instrumentation_config, "logger", logging.getLogger("codon_test_logger")) + + +def test_initialize_telemetry_uses_defaults(monkeypatch): + captured = {} + warnings = [] + + def capture_provider(provider): + captured["provider"] = provider + + def capture_warning(msg): + warnings.append(msg) + + _patch_base(monkeypatch, existing_provider=object()) + monkeypatch.setattr( + instrumentation_config.trace, + "set_tracer_provider", + capture_provider, + ) + monkeypatch.setattr(instrumentation_config.logger, "warning", lambda msg: warnings.append(msg)) + + instrumentation_config.initialize_telemetry() + + provider = captured["provider"] + assert provider.processors # processor added + exporter = provider.processors[0].span_exporter + assert exporter.endpoint == DEFAULT_INGEST_ENDPOINT + assert exporter.headers == {} + assert provider.resource.attributes["service.name"] == "unknown_codon_service" + assert warnings # warning about missing API key + + +def test_initialize_telemetry_prefers_arguments(monkeypatch): + monkeypatch.setenv("CODON_API_KEY", "env-key") + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://env:4317") + monkeypatch.setenv("OTEL_SERVICE_NAME", "env-service") + + captured = {} + _patch_base(monkeypatch, existing_provider=object()) + monkeypatch.setattr( + instrumentation_config.trace, + "set_tracer_provider", + lambda provider: captured.setdefault("provider", provider), + ) + + instrumentation_config.initialize_telemetry( + api_key="arg-key", + service_name="arg-service", + endpoint="http://arg:4317", + ) + + provider = captured["provider"] + exporter = provider.processors[0].span_exporter + assert exporter.endpoint == "http://arg:4317" + assert exporter.headers == {"x-codon-api-key": "arg-key"} + assert provider.resource.attributes["service.name"] == "arg-service" + + +def test_initialize_telemetry_env_fallback(monkeypatch): + monkeypatch.setenv("CODON_API_KEY", "env-key") + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://env:4317") + monkeypatch.setenv("OTEL_SERVICE_NAME", "env-service") + + captured = {} + _patch_base(monkeypatch, existing_provider=object()) + monkeypatch.setattr( + instrumentation_config.trace, + "set_tracer_provider", + lambda provider: captured.setdefault("provider", provider), + ) + + instrumentation_config.initialize_telemetry() + + provider = captured["provider"] + exporter = provider.processors[0].span_exporter + assert exporter.endpoint == "http://env:4317" + assert exporter.headers == {"x-codon-api-key": "env-key"} + assert provider.resource.attributes["service.name"] == "env-service" + + +def test_initialize_attach_to_existing_provider(monkeypatch): + existing = DummyProvider(resource=DummyResource(attributes={"service.name": "existing"})) + + _patch_base(monkeypatch, existing_provider=existing) + monkeypatch.setenv("CODON_ATTACH_TO_EXISTING_OTEL_PROVIDER", "true") + + set_called = {"called": False} + monkeypatch.setattr( + instrumentation_config.trace, + "set_tracer_provider", + lambda provider: set_called.__setitem__("called", True), + ) + + instrumentation_config.initialize_telemetry(api_key="attach-key", endpoint="http://arg:4317") + assert set_called["called"] is False + assert len(existing.processors) == 1 + exporter = existing.processors[0].span_exporter + assert exporter.headers == {"x-codon-api-key": "attach-key"} + assert exporter.endpoint == "http://arg:4317" + + # Calling again should not duplicate processors + instrumentation_config.initialize_telemetry(api_key="attach-key", endpoint="http://arg:4317", attach_to_existing=True) + assert len(existing.processors) == 1 + + +def test_attach_flag_argument_overrides_env(monkeypatch): + existing = DummyProvider(resource=DummyResource(attributes={})) + _patch_base(monkeypatch, existing_provider=existing) + monkeypatch.setenv("CODON_ATTACH_TO_EXISTING_OTEL_PROVIDER", "true") + + captured = {} + monkeypatch.setattr( + instrumentation_config.trace, + "set_tracer_provider", + lambda provider: captured.setdefault("provider", provider), + ) + + instrumentation_config.initialize_telemetry(attach_to_existing=False) + # attach disabled -> new provider set + assert "provider" in captured + assert captured["provider"] is not existing