[RFC] SkyRL TITO Proxy

[xinze-zheng]

Summary

This RFC focuses on the agent integration item in roadmap #1846.

A generator-owned proxy server exposes an OpenAI-compatible Chat Completions
endpoint to agent frameworks such as Harbor. Instead of forwarding chat messages
directly to vLLM, the proxy renders messages into token IDs, sends exact prompt
IDs to SkyRL's inference engine, records sampled completion IDs and logprobs, and
returns a compatible chat response to the agent.

For subsequent turns, the proxy preserves sampled token IDs and safely appends
new environment messages using a model-aware renderer. The recorded trace is
finally converted to GeneratorOutput.

Functional Requirements

1. Provide a /chat/completions-compatible endpoint that requires Harbor to
   change only its endpoint URL.
2. Transparently run token-in/token-out inference through the SkyRL router and
   record tokens, masks, logprobs, and potentially R3 data efficiently.
3. Handle compaction and expose branch and transition views.

Design

Trace Design

Trace is the bookkeeping object for one trial attempt.

• Trace maintains a message graph. Each node stores one canonical message, its
  parent ID, and the token, mask, and logprob delta introduced by that message.
• Trace finds the longest message prefix with an in-memory
  (parent_id, message_hash) -> node_id index.
• Before reusing a prefix, commit verifies that node token deltas exactly match
  the prompt IDs sent to inference.
• Every root-to-leaf path is one exact trajectory. Shared prefixes are stored
  once.
• Trace exposes branch and transition views without duplicating the underlying
  message graph.
• Trace's underlying data structure is similar to Trie, Therefore the complexity for matching is $O(m)$ where $m$ is the number of messages. We do hope the branch will not happen often so we will also maintain the hash of the latest context for fast match.

Proxy Design

• HarborGenerator.generate() starts one FastAPI proxy. Concurrent Harbor
  trials share that proxy.
• Each trial attempt creates a Trace and registers it with the proxy. Harbor's
  model endpoint is changed to the proxy endpoint.
• For the first call, the proxy renders the complete prompt. For later calls,
  Trace finds the latest matching message node and returns the new message tail.
• The renderer safely extends the previous exact (prompt_ids, completion_ids)
  with the new tail.
• The proxy runs token-in inference, commits the exact result to Trace, and then
  returns the Chat Completions response to the agent.

Example

The easiest way to illustrate this design is the following example.

Round 1

• Proxy receives initial messages $R_1$ and renders prompt IDs $P_1$.
• Inference returns completion IDs $C_1$ and logprobs $L_1$. Renderer parses
  $C_1$ into assistant message $M_1$.
• Trace commits one node per message in $R_1$, followed by a sampled $M_1$ node.
• Proxy returns $M_1$ to the agent.

Round 2

• Agent runs an environment step and receives message $T_1$.
• Agent sends $R_1 \Vert [M_1, T_1]$.
• Trace matches $R_1 \Vert [M_1]$, recovers exact $(P_1, C_1)$, and identifies
  $[T_1]$ as the new tail.
• Renderer constructs $P_2 = \mathrm{Bridge}(P_1, C_1, [T_1])$.
• Inference returns $(C_2, L_2)$ and renderer parses assistant message $M_2$.
• Trace commits a nonsampled $T_1$ node and a sampled $M_2$ node. The path
  reconstructs $P_2 \Vert C_2$.
• Proxy returns $M_2$ to the agent.

Round 3

• Agent sends $R_1 \Vert [M_1, T_1, M_2, T_2]$.
• Trace matches through $M_2$ and identifies $[T_2]$ as the new tail.
• Renderer constructs $P_3 = \mathrm{Bridge}(P_2, C_2, [T_2])$.
• Inference returns $(C_3, L_3)$ and renderer parses $M_3$.
• Trace commits $T_2$ and $M_3$. The path reconstructs $P_3 \Vert C_3$.
• Proxy returns $M_3$ to the agent.

Branch Cases

• Linear extension: Message and token prefixes match. New nodes extend the
  current leaf.
• Compaction or rewritten history: Matching stops at the last unchanged
  message. New messages form a branch from that node or the dummy root.
• Unsafe renderer extension: Renderer returns None. The proxy fully renders
  the prompt, and exact token comparison selects the longest reusable node.
• Retry: An identical retry reuses the in-flight or cached response and does
  not commit new nodes.

[CharlieFRuan]

Thanks a lot for the RFC! Some comments below

compatible endpoint that requires Harbor to change only its endpoint URL. HarborGenerator.generate() starts one FastAPI proxy. Concurrent Harbor trials share that proxy.

Should work with arbitrary GeneratorInterface implementation (e.g. MiniSWEAgent), but for the sake of RFC we can center around HarborGenerator.

Before reusing a prefix, commit verifies that node token deltas exactly match the prompt IDs sent to inference.

For prefix matching in Trace, why do we need it to be in token space?

and then returns the Chat Completions response to the agent.

This would require translating vLLM ChatCompletionsRequests to CompletionsRequests, and CompletionResponse back to ChatCompletionResponse. I'd imagine there cannot be one-to-one mapping for all the kwargs for this translation so we should have unit tests / assertions on fields that might not be translatable. Just wanted to call this out.

Re Example

The usage of [T1] rather than T1 -- does it suggest that the environment step can produce more than one message? If so, should each become a node in Trace, or should it be a single node? I don't have a clear answer, but it depends on how prefix matching works. If we have [T1, T2] in a single turn, and request 3 from the user starts with P1, C1, T1 (i.e. drops T2) -- do we want to match that T1 message/node.

Inference returns $(C_2, L_2)$ and renderer parses assistant message $M_2$.

What is $L_2$ here?

Unsafe renderer extension: Renderer returns None

Could you elaborate when renderer returns None?

The first concrete goal can be:

• For HarborGenerator, we currently perform step-wise training with terminus 2, where we merge ourselves post-hoc and it cannot always merge (since we re-tokenize in Harbor Terminus 2's multi-turn requests) See [train][step-wise] Add prefix-aware merging for step-wise training #1538
• So an initial goal is to replace HarborGenerator's usage of step-wise with the proxy, and we can always successfully merge (i.e. each trajectory is precisely a single sequence)

High-level comment

Would be great to see what python classes will need to be introduced, and how that class interacts with a GeneratorInterface implementation, and what API would the proxy/Trace expose for the user to re-compose to GeneratorOutput. I'd imagine the trace_to_generator_output() can be highly customizable

[xinze-zheng]

Thank you for the feedback.

For prefix matching in Trace, why do we need it to be in token space?

For Trace prefix matching, only message space is enough for most cases. The only case we need token level matching is when a full re-render happens. In this case we still want to use the message graph to avoid re-store the same node. However, full re-render might introduce retokenization, in this case, only message space match will be incorrect.

Could you elaborate when renderer returns None?

TL;DR
When thinking retention policy is all, and previous context is not malformatted, it will never return None.

Here we are referring to the renderer.bridge_to_next_turn. Concretely for Qwen3, in the following cases the renderer will return None.

1. new_messages contain a assistant turn. Essentially refusing to retokenize a completion response.
2. Thinking retention policy. If the policy uses chat-template and remove the thinking tokens, then full-render is needed. For Qwen3, if thinking is enabled, the policy is tool_cycle, meaning renderer will try to bridge tool cycles and re-render at new user-query boundary. For SkyRL, we could always retain thinking tokens with all policy.
3. role is unknown (not user, tool, system)
4. bridge has not valid input arguments (this won't happen if using correctly)

What is $L_2$ here?

It's the logprob of the completion.

The usage of [T1] rather than T1 -- does it suggest that the environment step can produce more than one message?

Yes. Technically there can be multiple tool responses or even follow by a user query. In current harbor example it's always a single tool response. To me (so is the design of verifer), the more fine-grained the better. Making every message a node will allow maximum usage of nodes. In you example if user drops T2, then we can still match T1.

Python classes proposal

Proxy and Trace

@dataclass
class MessageNode:
  parent_id: int | None
  message: Message
  token_ids: list[int]
  sampled_mask: list[bool]
  logprobs: list[float]

@dataclass(frozen=True)
class PendingTurn:
    messages: list[Message]
    prefix_node_ids: list[int]
    tail: list[Message]

    def bridge_anchor(
        self,
    ) -> tuple[list[int], list[int]] | None:
        """Return previous exact prompt and completion IDs for renderer."""

class Trace:
  nodes: list[MessageNode]

  def prepare_turn(self, messages: list[Message]) -> PendingTurn: ...
  def commit(self, turn: PendingTurn, result: ModelTurnResult) -> CommitResult: ...
  #Optionally provide views for user
  def branches(self) -> list[BranchView]: ...
  def transitions(self, branch_id: int) -> list[TransitionView]: ...


@dataclass(frozen=True)
class ProxySlot:
  base_url: str
  secret: str

@dataclass
class TITOProxyConfig:
    host: str = "127.0.0.1"
    port: int = 0
    max_request_bytes: int = 64 * 1024 * 1024

@dataclass(frozen=True)
class ProxyHandle:
    base_url: str
    session_id: str


class TITOProxy:
    @asynccontextmanager
    async def register(
        self,
        trace: Trace,
        *,
        session_id: str,
        cache_salt: str | None,
    ) -> AsyncIterator[ProxyHandle]:
        ...

Conversion to SkyRL training data.

@dataclass(frozen=True)
class TraceOutcome:
  trace: Trace
  trajectory_id: TrajectoryID
  reward: float
  stop_reason: str
  generation_time: float | None


class TraceComposer(Protocol):
  def compose(self, outcomes: list[TraceOutcome]) -> GeneratorOutput: ...

Generator integration. Prepare trace, register to proxy, finally call composer.

class HarborGenerator(GeneratorInterface):
  def __init__(self, ..., composer: TraceComposer):
    self.proxy = TITOProxy(inference_engine_client, renderer)
    self.composer = composer

  async def generate(self, input_batch: GeneratorInput) -> GeneratorOutput:
    outcomes = []
    async with self.proxy.serving():
      outcomes = await self._run_trials(input_batch)
    return self.composer.compose(outcomes)

  async def _run_attempt(self, trajectory_id, prompt, cache_salt):
    trace = Trace()
    session_id = uuid4().hex

    async with self.proxy.register(
        trace,
        session_id=session_id,
        cache_salt=cache_salt,
    ) as handle:
        result = await run_harbor_trial(
            api_base=f"{handle.base_url}/v1",
            session_id=handle.session_id,
        )
    return TraceOutcome(
      trace=trace,
      trajectory_id=trajectory_id,
      reward=result.reward,
      stop_reason=result.stop_reason,
      generation_time=result.generation_time,
    )

[CharlieFRuan]

Appreciate the detailed response! I think the design makes sense overall.

I think some class definitions are missing, e.g. ModelTurnResult, TransitionView, BranchView . And how do we define a branch?

Some descriptions / pseudocode for Trace class methods will be helpful. E.g. which method is calling which method, and when. Also when / who is calling bridge_anchor.

We can test the following cases:

1. The most well behaved all-linearly-appending case, where we should expect a single training sequence
2. When we perform summarization once, then we should expect two training sequences
3. Case 1 but we strip thinking tokens
4. When the model re-tries a request (e.g. due to timeout)

Some diagrams could be helpful for examples 2, 3, and 4.

[xinze-zheng]

I have updated a more detailed design doc and implemented the corresponding draft. The doc will cover, code flow, more detailed class and their APIs, and examples.