fix gpt oss tool calling

dashboard/src/lib/stores/app.svelte.ts

@@ -72,8 +72,6 @@ export interface Instance {
     runnerToShard?: Record<string, unknown>;
     nodeToRunner?: Record<string, string>;
   };
-  draftModel?: string;
-  numDraftTokens?: number;
 }
 
 // Granular node state types from the new state structure

src/exo/master/main.py

@@ -24,7 +24,6 @@ from exo.shared.types.commands import (
     PlaceInstance,
     RequestEventLog,
     SendInputChunk,
-    SetInstanceDraftModel,
     TaskFinished,
     TestCommand,
     TextGeneration,
@@ -36,7 +35,6 @@ from exo.shared.types.events import (
     IndexedEvent,
     InputChunkReceived,
     InstanceDeleted,
-    InstanceDraftModelUpdated,
     NodeGatheredInfo,
     NodeTimedOut,
     TaskCreated,
@@ -321,14 +319,6 @@ class Master:
                                     chunk=chunk,
                                 )
                             )
-                        case SetInstanceDraftModel():
-                            generated_events.append(
-                                InstanceDraftModelUpdated(
-                                    instance_id=command.instance_id,
-                                    draft_model=command.draft_model,
-                                    num_draft_tokens=command.num_draft_tokens,
-                                )
-                            )
                         case TaskFinished():
                             generated_events.append(
                                 TaskDeleted(

src/exo/master/placement.py

@@ -153,8 +153,6 @@ def place_instance(
                 shard_assignments=shard_assignments,
                 jaccl_devices=mlx_jaccl_devices,
                 jaccl_coordinators=mlx_jaccl_coordinators,
-                draft_model=command.draft_model,
-                num_draft_tokens=command.num_draft_tokens,
             )
         case InstanceMeta.MlxRing:
             ephemeral_port = random_ephemeral_port()
@@ -169,8 +167,6 @@ def place_instance(
                 shard_assignments=shard_assignments,
                 hosts_by_node=hosts_by_node,
                 ephemeral_port=ephemeral_port,
-                draft_model=command.draft_model,
-                num_draft_tokens=command.num_draft_tokens,
             )
 
     return target_instances

src/exo/shared/apply.py

@@ -12,7 +12,6 @@ from exo.shared.types.events import (
     InputChunkReceived,
     InstanceCreated,
     InstanceDeleted,
-    InstanceDraftModelUpdated,
     NodeDownloadProgress,
     NodeGatheredInfo,
     NodeTimedOut,
@@ -73,8 +72,6 @@ def event_apply(event: Event, state: State) -> State:
             return apply_instance_created(event, state)
         case InstanceDeleted():
             return apply_instance_deleted(event, state)
-        case InstanceDraftModelUpdated():
-            return apply_instance_draft_model_updated(event, state)
         case NodeTimedOut():
             return apply_node_timed_out(event, state)
         case NodeDownloadProgress():
@@ -193,25 +190,6 @@ def apply_instance_deleted(event: InstanceDeleted, state: State) -> State:
     return state.model_copy(update={"instances": new_instances})
 
 
-def apply_instance_draft_model_updated(
-    event: InstanceDraftModelUpdated, state: State
-) -> State:
-    if event.instance_id not in state.instances:
-        return state
-    instance = state.instances[event.instance_id]
-    updated_instance = instance.model_copy(
-        update={
-            "draft_model": event.draft_model,
-            "num_draft_tokens": event.num_draft_tokens,
-        }
-    )
-    new_instances: Mapping[InstanceId, Instance] = {
-        **state.instances,
-        event.instance_id: updated_instance,
-    }
-    return state.model_copy(update={"instances": new_instances})
-
-
 def apply_runner_status_updated(event: RunnerStatusUpdated, state: State) -> State:
     new_runners: Mapping[RunnerId, RunnerStatus] = {
         **state.runners,

src/exo/shared/types/commands.py

@@ -38,8 +38,6 @@ class PlaceInstance(BaseCommand):
     sharding: Sharding
     instance_meta: InstanceMeta
     min_nodes: int
-    draft_model: ModelId | None = None
-    num_draft_tokens: int = 4
 
 
 class CreateInstance(BaseCommand):
@@ -74,14 +72,6 @@ class DeleteDownload(BaseCommand):
     model_id: ModelId
 
 
-class SetInstanceDraftModel(BaseCommand):
-    """Set or update the draft model for an existing instance."""
-
-    instance_id: InstanceId
-    draft_model: ModelId | None  # None to disable speculative decoding
-    num_draft_tokens: int = 4
-
-
 class CancelDownload(BaseCommand):
     target_node_id: NodeId
     model_id: ModelId
@@ -99,7 +89,6 @@ Command = (
     | PlaceInstance
     | CreateInstance
     | DeleteInstance
-    | SetInstanceDraftModel
     | TaskFinished
     | SendInputChunk
 )

src/exo/shared/types/events.py

@@ -5,7 +5,7 @@ from pydantic import Field
 
 from exo.shared.topology import Connection
 from exo.shared.types.chunks import GenerationChunk, InputImageChunk
-from exo.shared.types.common import CommandId, Id, ModelId, NodeId, SessionId
+from exo.shared.types.common import CommandId, Id, NodeId, SessionId
 from exo.shared.types.tasks import Task, TaskId, TaskStatus
 from exo.shared.types.worker.downloads import DownloadProgress
 from exo.shared.types.worker.instances import Instance, InstanceId
@@ -68,14 +68,6 @@ class InstanceDeleted(BaseEvent):
     instance_id: InstanceId
 
 
-class InstanceDraftModelUpdated(BaseEvent):
-    """Draft model updated on an existing instance."""
-
-    instance_id: InstanceId
-    draft_model: ModelId | None
-    num_draft_tokens: int
-
-
 class RunnerStatusUpdated(BaseEvent):
     runner_id: RunnerId
     runner_status: RunnerStatus
@@ -149,7 +141,6 @@ Event = (
     | TaskAcknowledged
     | InstanceCreated
     | InstanceDeleted
-    | InstanceDraftModelUpdated
     | RunnerStatusUpdated
     | RunnerDeleted
     | NodeTimedOut

src/exo/shared/types/tasks.py

@@ -40,12 +40,6 @@ class DownloadModel(BaseTask):  # emitted by Worker
     shard_metadata: ShardMetadata
 
 
-class DownloadDraftModel(BaseTask):  # emitted by Worker
-    """Download a draft model for speculative decoding (rank 0 only)."""
-
-    model_id: str  # HuggingFace model ID
-
-
 class LoadModel(BaseTask):  # emitted by Worker
     pass
 
@@ -86,17 +80,9 @@ class Shutdown(BaseTask):  # emitted by Worker
     runner_id: RunnerId
 
 
-class SetDraftModel(BaseTask):  # emitted by Worker
-    """Load or clear a draft model on an already-running instance."""
-
-    model_id: str | None  # HuggingFace model ID, or None to clear
-    num_draft_tokens: int = 4
-
-
 Task = (
     CreateRunner
     | DownloadModel
-    | DownloadDraftModel
     | ConnectToGroup
     | LoadModel
     | StartWarmup
@@ -104,5 +90,4 @@ Task = (
     | ImageGeneration
     | ImageEdits
     | Shutdown
-    | SetDraftModel
 )

src/exo/shared/types/worker/instances.py

@@ -2,7 +2,7 @@ from enum import Enum
 
 from pydantic import model_validator
 
-from exo.shared.types.common import Host, Id, ModelId, NodeId
+from exo.shared.types.common import Host, Id, NodeId
 from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
 from exo.utils.pydantic_ext import CamelCaseModel, TaggedModel
 
@@ -19,8 +19,6 @@ class InstanceMeta(str, Enum):
 class BaseInstance(TaggedModel):
     instance_id: InstanceId
     shard_assignments: ShardAssignments
-    draft_model: ModelId | None = None  # For speculative decoding (rank 0 only)
-    num_draft_tokens: int = 4  # Tokens to draft per iteration (when draft_model is set)
 
     def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
         return self.shard_assignments.runner_to_shard.get(runner_id, None)

src/exo/worker/engines/mlx/utils_mlx.py

@@ -223,27 +223,6 @@ def load_mlx_items(
     return cast(Model, model), tokenizer
 
 
-def load_draft_model(model_id: ModelId) -> nn.Module:
-    """Load a draft model for speculative decoding (rank 0 only).
-
-    Draft models are small models (typically 0.5B-2B parameters) used to
-    generate candidate tokens quickly, which are then verified by the main
-    model in a single forward pass.
-
-    Assumes the model has already been downloaded by the worker.
-
-    Args:
-        model_id: HuggingFace model ID for the draft model
-
-    Returns:
-        The loaded draft model
-    """
-    model_path = build_model_path(model_id)
-    draft_model, _ = load_model(model_path, strict=True)
-    logger.info(f"Loaded draft model from {model_path}")
-    return draft_model
-
-
 def shard_and_load(
     shard_metadata: ShardMetadata,
     group: Group,
@@ -337,6 +316,8 @@ def get_eos_token_ids_for_model(model_id: ModelId) -> list[int] | None:
         return [154820, 154827, 154829]
     elif "glm" in model_id_lower:
         return [151336, 151329, 151338]
+    elif "gpt-oss" in model_id_lower:
+        return [200002, 200012]
     return None
 
 

src/exo/worker/runner/runner.py

@@ -11,6 +11,7 @@ from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 from openai_harmony import (  # pyright: ignore[reportMissingTypeStubs]
     HarmonyEncodingName,
+    HarmonyError,  # pyright: ignore[reportUnknownVariableType]
     Role,
     StreamableParser,
     load_harmony_encoding,
@@ -568,7 +569,11 @@ def parse_gpt_oss(
 
     for response in responses:
         assert isinstance(response, GenerationResponse)
-        stream.process(response.token)
+        try:
+            stream.process(response.token)
+        except HarmonyError:
+            logger.error("Encountered critical Harmony Error, returning early")
+            return
 
         delta = stream.last_content_delta
         ch = stream.current_channel