Mosaic: Confidence-Aware Contrastive Learning over Multiplex Tool Graphs for Robust Task Planning

Task planning requires LLM-based agents to translate a complex user request into an ordered sequence of tool invocations. Graph-based task planners address this problem by learning over a tool graph that encodes which tools can be invoked in sequence, and consistently outperform planners that rely on LLMs alone. However, the tool graph is rarely available and is built heuristically, either from historical tool invocations or from input-output type compatibility; therefore, its edges vary widely in reliability, yet existing planners treat them as equally trustworthy binary connections and typically rely on a single construction. We propose Mosaic (**M**ultiplex c**O**nfidence-aware **S**tructure-**A**ugmented **I**dentity-aware **C**ontrastive learning), a model-agnostic plug-in method that makes the tool representations of graph-based planners robust to unreliable tool graphs. Mosaic (i) combines the historical and dependency graphs into a confidence-weighted multiplex tool graph; (ii) trains relation-specific graph encoders with a contrastive objective whose augmentations drop low-confidence edges more aggressively and which aligns each tool across the two relations; and (iii) fuses the resulting embeddings with learnable, graph-agnostic tool identity embeddings. Across three TaskBench domains, three LLMs, and two state-of-the-art planners, Mosaic consistently improves planning performance (with relative accuracy gains of up to 22.25%), remains more robust under structural noise injected into the tool graphs, and adds only marginal computational overhead.