query method

Implementation

AgentQuery query({required String prompt, Options? options}) {
  final streamId = _nextStreamId();
  final controller = StreamController<SdkMessage>();
  StreamSubscription<Map<String, dynamic>>? sub;
  var finished = false;

  // Index any Dart-defined (`sdk`) MCP tool handlers so the extension can
  // invoke them mid-query over the reverse RPC (todo #5).
  final toolRegistry = AgentSdkToolRegistry();
  final hasTools = toolRegistry.addServers(options?.mcpServers);

  // The caller's `canUseTool` approval callback, dispatched over the reverse
  // RPC whenever the model requests a tool (todo #6).
  final canUseTool = options?.canUseTool;

  Future<void> finish({bool cancelRemote = false}) async {
    if (finished) return;
    finished = true;
    if (hasTools) {
      transport.unregisterTools(streamId);
    }
    if (canUseTool != null) {
      transport.unregisterCanUseTool(streamId);
    }
    await sub?.cancel();
    if (cancelRemote) {
      try {
        await transport.cancelQuery(streamId);
      } catch (_) {
        // Best-effort: the remote may already have completed.
      }
    }
    if (!controller.isClosed) {
      await controller.close();
    }
  }

  controller.onListen = () {
    // Subscribe to chunks *before* starting the query so no early message
    // is dropped between start and subscription.
    sub = transport.chunks
        .where((chunk) => chunk['streamId'] == streamId)
        .listen((chunk) {
          if (finished) return;
          final error = chunk['error'];
          if (error != null) {
            controller.addError(AgentSdkQueryException(error.toString()));
            finish();
            return;
          }
          if (chunk['done'] == true) {
            finish();
            return;
          }
          final message = chunk['message'];
          if (message is Map) {
            controller.add(
              SdkMessage.fromJson(message.cast<String, dynamic>()),
            );
          }
        });

    // Register tool handlers before starting so an early `agentSdk.toolCall`
    // cannot arrive before the registry is in place.
    if (hasTools) {
      transport.registerTools(streamId, toolRegistry);
    }

    // Likewise register the approval callback before starting so an early
    // `agentSdk.canUseTool` request has a handler.
    if (canUseTool != null) {
      transport.registerCanUseTool(streamId, canUseTool);
    }

    transport
        .startQuery({
          'streamId': streamId,
          'prompt': prompt,
          if (options != null) 'options': options.toJson(),
        })
        .catchError((Object e) {
          if (finished) return;
          controller.addError(e);
          finish();
        });
  };

  // A subscriber cancelling the stream aborts the underlying query.
  controller.onCancel = () => finish(cancelRemote: true);

  return AgentQuery._(
    controller.stream,
    interrupt: () => finish(cancelRemote: true),
  );
}