Restate tracing with Future AGI

pip install 'restate-sdk==1.0.5' hypercorn fi-instrumentation-otel traceai-openai openai

Install Restate server 1.8.0 or 1.7.13 from its release assets. Follow Restate’s quickstart for the CLI. Keep the server on one of these versions for this recipe; the collector command below pins OpenTelemetry Collector contrib to 0.162.0.

Save this as otelcol.yaml:

# OpenTelemetry Collector (contrib) between Restate and Future AGI.
#
# Restate's server and your Restate service both export to this collector. It then
#   1. names the Observe project (resource attribute project_name) when the sender did not,
#   2. removes host names, network addresses and stack traces, always,
#   3. removes error text and Virtual Object / Workflow keys, unless KEEP_RESTATE_DETAILS=true,
#   4. forwards plain (uncompressed) OTLP/HTTP with the X-Api-Key and X-Secret-Key headers
#      that each sender put on its own request. This collector holds no Future AGI keys.
#
# It listens on /tracer/v1/traces, the same path Future AGI uses, so traceAI's register()
# can be pointed here with FI_BASE_URL=http://<this host>:4318.
#
# Environment:
#   FI_TRACES_URL         where to forward, e.g. https://api.futureagi.com/tracer/v1/traces
#   FI_PROJECT_NAME       Observe project for spans that arrive without project_name
#   KEEP_RESTATE_DETAILS  "true" keeps error text and Virtual Object / Workflow keys (default: removed)

extensions:
  # Copies the two Future AGI key headers from the incoming request onto the outgoing one.
  headers_setter:
    headers:
      - action: upsert
        key: X-Api-Key
        from_context: X-Api-Key
      - action: upsert
        key: X-Secret-Key
        from_context: X-Secret-Key

receivers:
  otlp:
    protocols:
      http:
        endpoint: 0.0.0.0:4318
        traces_url_path: /tracer/v1/traces
        include_metadata: true   # needed for headers_setter to see the incoming headers

processors:
  # `insert` keeps a project_name the sender already set (traceAI's register() sets one).
  resource/future_agi:
    attributes:
      - key: project_name
        value: ${env:FI_PROJECT_NAME}
        action: insert
      - key: project_type
        value: observe
        action: insert

  # Always removed. service.instance.id is "<cluster>/<node-name>", and Restate's node name
  # defaults to the machine's host name. Stack traces carry local paths and user names.
  transform/always:
    error_mode: ignore
    trace_statements:
      - context: resource
        statements:
          - delete_key(attributes, "service.instance.id")
          - delete_key(attributes, "host.name")
      - context: span
        statements:
          - delete_key(attributes, "restate.deployment.address")
          - delete_key(attributes, "client.socket.address")
          - delete_key(attributes, "client.socket.port")
          - delete_key(attributes, "client.address")
          - delete_key(attributes, "client.port")
      - context: spanevent
        statements:
          - delete_key(attributes, "exception.stacktrace")

  # Removed unless KEEP_RESTATE_DETAILS=true. Error text can quote user data, and a
  # Virtual Object key is often a user id or an e-mail address. Restate already writes
  # span names with {key}; this does the same for restate.invocation.target.
  transform/details:
    error_mode: ignore
    trace_statements:
      - context: span
        conditions:
          - '"${env:KEEP_RESTATE_DETAILS:-false}" != "true"'
        statements:
          - delete_key(attributes, "error.message")
          - set(status.message, "error") where status.message != ""
          - replace_pattern(attributes["restate.invocation.target"], "^([^/]+)/.+/([^/]+)$$", "$$1/{key}/$$2")
      - context: spanevent
        conditions:
          - '"${env:KEEP_RESTATE_DETAILS:-false}" != "true"'
        statements:
          - delete_key(attributes, "exception.message")

exporters:
  otlp_http/future_agi:
    traces_endpoint: ${env:FI_TRACES_URL}
    encoding: proto
    compression: none          # optional: Future AGI cloud and self-hosted v1.56.0+ also accept gzip
    auth:
      authenticator: headers_setter

service:
  extensions: [headers_setter]
  telemetry:
    metrics:
      level: none              # no self-metrics endpoint; this collector only carries traces
    logs:
      level: warn
  pipelines:
    traces:
      receivers: [otlp]
      processors: [resource/future_agi, transform/always, transform/details]
      exporters: [otlp_http/future_agi]

Start it:

docker run -d --name restate-otelcol -p 127.0.0.1:4318:4318 \
  -v "$PWD/otelcol.yaml:/etc/otelcol.yaml:ro" \
  -e FI_TRACES_URL=https://api.futureagi.com/tracer/v1/traces \
  -e FI_PROJECT_NAME=restate-agent \
  otel/opentelemetry-collector-contrib:0.162.0 --config /etc/otelcol.yaml

For a self-hosted Future AGI, set FI_TRACES_URL to your instance’s /tracer/v1/traces URL. The collector holds no keys: each sender’s X-Api-Key and X-Secret-Key headers are copied onto the request it forwards. Keep port 4318 on a private network.

Save this as restate.toml and start the server with restate-server --config-file restate.toml:

# Restate server: export invocation traces to the OpenTelemetry Collector in otelcol.yaml.
#
#   restate-server --config-file restate.toml
#
# otlp+http:// selects OTLP/HTTP (protobuf). A plain http:// endpoint would mean OTLP/gRPC.
# The path is the collector's receiver path in otelcol.yaml.
tracing-endpoint = "otlp+http://localhost:4318/tracer/v1/traces"

# Sent on every export request. The collector copies them onto the request it makes to
# Future AGI. Rather than writing keys into this file you can set them in the environment:
#   RESTATE_TRACING_HEADERS='{X-Api-Key="<FI_API_KEY>", X-Secret-Key="<FI_SECRET_KEY>"}'
[tracing-headers]
X-Api-Key = "<FI_API_KEY>"
X-Secret-Key = "<FI_SECRET_KEY>"

Both headers are required. Restate’s docs show a single authorization header as the example; that header does not authenticate to Future AGI. Restate sends the export uncompressed, which is what Future AGI accepts.

The environment form keeps the keys out of the file. Any server setting can be overridden this way (RESTATE_TRACING_ENDPOINT, for example). In Docker, use host.docker.internal or the collector’s service name instead of localhost.

The snippets below show the tracing setup in the complete agent_service.py example; they are not a standalone service. The four lines that matter are the register() call, the RestateTracerProvider wrapper, instrument(tracer_provider=...) with that wrapper, and FITracer over the wrapper for your own spans.

import os

import restate
from fi_instrumentation import FITracer, TraceConfig, register, using_session
from fi_instrumentation.fi_types import ProjectType
from openai import AsyncOpenAI
from restate.ext.tracing import RestateTracerProvider
from traceai_openai import OpenAIInstrumentor

# 1. A Future AGI tracer provider. Not registered globally: only the Restate-aware wrapper below uses it.
tracer_provider = register(
    project_name=os.environ.get("FI_PROJECT_NAME", "restate-agent"),
    project_type=ProjectType.OBSERVE,
    set_global_tracer_provider=False,
    verbose=False,
)
# 2. Every span created through this wrapper becomes a child of the Restate invocation attempt
#    that is running, and no span is created while Restate replays already-journaled steps.
restate_tracer_provider = RestateTracerProvider(tracer_provider)
OpenAIInstrumentor().instrument(tracer_provider=restate_tracer_provider)
tracer = FITracer(restate_tracer_provider.get_tracer(__name__), config=TraceConfig())

# Restate retries failed steps itself, so the OpenAI client should not retry on its own.
client = AsyncOpenAI(max_retries=0)

Put every model call and tool call in its own ctx.run step, so a completed step is journaled and never repeated:

agent = restate.VirtualObject("WeatherAgent")  # keyed by user: one conversation per user at a time


@agent.handler()
async def ask(ctx: restate.ObjectContext, req: dict) -> dict:
    invocation_id = ctx.request().id
    model = req.get("model", MODEL)
    messages = [
        {"role": "system", "content": "Use the get_weather tool, then answer in one sentence."},
        {"role": "user", "content": req["question"]},
    ]
    # No application session id? Use the Restate invocation id, so one invocation is one session.
    with using_session(req.get("session_id") or invocation_id):
        try:
            plan = await ctx.run_typed(
                "llm-plan", chat, restate.RunOptions(max_attempts=2), model=model, messages=messages, with_tools=True
            )
        except TerminalError:
            raise TerminalError("the model call failed")
        messages.append(plan)
        for call in plan["tool_calls"]:
            result = await ctx.run_typed(
                "tool-get_weather",
                get_weather,
                invocation_id=invocation_id,
                arguments=call["function"]["arguments"],
                flaky=bool(req.get("flaky")),
            )
            messages.append({"role": "tool", "tool_call_id": call["id"], "content": result})
        final = await ctx.run_typed("llm-answer", chat, model=model, messages=messages, with_tools=False)
    return {"answer": final["content"]}

The tool span is created with record_exception=False and keeps only the exception type as its status, so a failed tool step exports no message text and no stack trace (see create tool spans for the span helpers):

    with tracer.start_as_current_span(
        "get_weather", fi_span_kind="tool", record_exception=False, set_status_on_exception=False
    ) as span:

Run the service with traceAI pointed at the collector, so the collector can also strip error text from your spans:

export FI_API_KEY=... FI_SECRET_KEY=... FI_PROJECT_NAME=restate-agent
export FI_BASE_URL=http://localhost:4318      # the collector, not Future AGI
export OPENAI_API_KEY=...                     # and OPENAI_BASE_URL for an OpenAI-compatible endpoint
python agent_service.py                       # serves on 127.0.0.1:9080

curl localhost:9070/deployments --json '{"uri": "http://localhost:9080"}'   # or: restate deployments register http://localhost:9080
curl localhost:8080/WeatherAgent/user-1/ask -H 'content-type: application/json' \
  -d '{"question": "What is the weather in Lisbon today?"}'

The trace shows up in Observe under the project from FI_PROJECT_NAME. To see a retry, send "flaky": true: the first weather lookup fails and Restate retries the attempt.