use anyhow::{Context, Result};
use rmcp::{
handler::server::wrapper::Parameters,
model::{CallToolResult, ContentBlock as Content},
schemars, tool, tool_router, ErrorData as McpError,
};
use serde::{Deserialize, Serialize};
use crate::cli::drive::docs::info::InfoOutcome;
use crate::cli::drive::helpers::create_client_for;
use crate::drive::client::DriveClient;
use crate::drive::docs::api::{DocsApi, SuggestionsViewMode};
use crate::drive::docs::client::DocsClient;
use crate::drive::docs::read::{read, ReadOptions};
use crate::drive::docs::target;
use crate::drive::files_api::FilesApi;
use crate::mcp::drive_tools::account_param_doc;
use super::error::tool_error;
use super::git_tools::build_truncated_result;
use super::output_file;
use super::server::OmniDevServer;
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct DriveDocsInfoParams {
pub document_id: String,
#[doc = account_param_doc!()]
#[serde(default)]
pub account: Option<String>,
}
#[derive(Debug, Deserialize, schemars::JsonSchema)]
pub struct DriveDocsReadParams {
pub document_id: String,
#[serde(default)]
pub tab: Option<String>,
#[serde(default)]
pub suggestions_view: Option<String>,
#[serde(default)]
pub output_file: Option<String>,
#[doc = account_param_doc!()]
#[serde(default)]
pub account: Option<String>,
}
#[allow(missing_docs)] #[tool_router(router = drive_docs_tool_router, vis = "pub")]
impl OmniDevServer {
#[tool(
description = "Show a document's title, revision id and structural outline: named \
ranges and, per tab, heading text with its `[start,end)` index range plus \
paragraph/table/section-break counts — no full body text. `revision_id` is \
what a later `drive docs replace`/`append` presents as \
`writeControl.requiredRevisionId`, so a write against a document that moved \
underneath it is refused rather than misapplied; it is absent when the \
caller lacks edit access. Use `drive_docs_read` for the full element list \
with every index. \
Read-only. Mirrors `omni-dev drive docs info`. Output is YAML."
)]
pub async fn drive_docs_info(
&self,
Parameters(params): Parameters<DriveDocsInfoParams>,
) -> Result<CallToolResult, McpError> {
let client = create_client_for(params.account.as_deref()).map_err(tool_error)?;
let docs = DocsClient::from_drive_client(&client).map_err(tool_error)?;
let yaml = run_docs_info(&client, &docs, ¶ms)
.await
.map_err(tool_error)?;
Ok(build_truncated_result(yaml))
}
#[tool(
description = "Read a document's structural elements — paragraphs, tables, section \
breaks and more — each with its `[start,end)` index range in \
UTF-16 code units, plus headers/footers/footnotes and the document's \
`revision_id`. This is the *model* channel (indices for a later edit); \
`drive_file_read`'s markdown export is the *prose* channel and has no way \
back to an index. `tab` restricts to one tab id (from `drive_docs_info`); \
omit to read every tab. `suggestions_view` controls which view the indices \
are reported against — indices from one view only apply to an edit made \
under that same view. When `output_file` is set, writes the YAML result to \
that path and returns a short summary instead — recommended for a large \
document. \
Read-only — no write gate, lease or dry-run applies (unlike `docs \
replace`/`append`, exposed by separate gated write tools). \
Mirrors `omni-dev drive docs read`. Output is YAML."
)]
pub async fn drive_docs_read(
&self,
Parameters(params): Parameters<DriveDocsReadParams>,
) -> Result<CallToolResult, McpError> {
let client = create_client_for(params.account.as_deref()).map_err(tool_error)?;
let docs = DocsClient::from_drive_client(&client).map_err(tool_error)?;
let wrote_to_file = params.output_file.is_some();
let text = run_docs_read(&client, &docs, ¶ms)
.await
.map_err(tool_error)?;
if wrote_to_file {
Ok(CallToolResult::success(vec![Content::text(text)]))
} else {
Ok(build_truncated_result(text))
}
}
}
async fn run_docs_info(
drive: &DriveClient,
docs: &DocsClient,
params: &DriveDocsInfoParams,
) -> Result<String> {
let document = match DocsApi::new(docs)
.get_document(¶ms.document_id, SuggestionsViewMode::default())
.await
{
Ok(document) => document,
Err(err) => {
return Err(target::explain_failure(
&FilesApi::new(drive),
¶ms.document_id,
"info",
err,
)
.await)
}
};
let outcome = InfoOutcome::of(¶ms.document_id, &document);
yaml_result(&outcome)
}
async fn run_docs_read(
drive: &DriveClient,
docs: &DocsClient,
params: &DriveDocsReadParams,
) -> Result<String> {
let suggestions = parse_suggestions_view(params.suggestions_view.as_deref())?;
let opts = ReadOptions {
document_id: params.document_id.clone(),
tab: params.tab.clone(),
suggestions,
};
let outcome = match read(&DocsApi::new(docs), &opts).await {
Ok(outcome) => outcome,
Err(err) => {
return Err(target::explain_failure(
&FilesApi::new(drive),
¶ms.document_id,
"read",
err,
)
.await)
}
};
let yaml = yaml_result(&outcome)?;
match params.output_file.as_deref() {
Some(path) => output_file::write_to_file_yaml(path, &yaml, "yaml"),
None => Ok(yaml),
}
}
fn parse_suggestions_view(raw: Option<&str>) -> Result<SuggestionsViewMode> {
match raw.map(str::to_ascii_lowercase).as_deref() {
None | Some("default") => Ok(SuggestionsViewMode::DefaultForCurrentAccess),
Some("inline") => Ok(SuggestionsViewMode::Inline),
Some("accepted") => Ok(SuggestionsViewMode::PreviewAccepted),
Some("without") => Ok(SuggestionsViewMode::PreviewWithoutSuggestions),
Some(other) => {
anyhow::bail!(
"unknown suggestions_view {other:?} (expected 'default', 'inline', 'accepted', \
or 'without')"
)
}
}
}
fn yaml_result<T: Serialize>(data: &T) -> Result<String> {
serde_yaml::to_string(data).context("Failed to serialize result as YAML")
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use rmcp::handler::server::wrapper::Parameters;
use wiremock::matchers::{method, path, query_param};
use wiremock::{Mock, MockServer, ResponseTemplate};
use super::*;
use crate::drive::auth::{DriveCredentials, DriveGrantedScopes};
use crate::drive::docs::client::DOCS_API_URL;
use crate::drive::test_support::EnvGuard;
use crate::test_support::env::MapEnv;
use crate::utils::secret::Secret;
fn test_credentials() -> DriveCredentials {
DriveCredentials {
client_id: "client-1".to_string(),
client_secret: Secret::new("secret-1"),
refresh_token: Secret::new("refresh-1"),
scope: DriveGrantedScopes::READONLY,
}
}
async fn docs_and_drive_clients(server: &MockServer) -> (DriveClient, DocsClient) {
Mock::given(method("POST"))
.and(path("/token"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"access_token": "test-token",
"expires_in": 3600,
})))
.mount(server)
.await;
let mut drive = DriveClient::new(&server.uri(), &test_credentials()).unwrap();
crate::drive::client::test_support::replace_session(
&mut drive,
&test_credentials(),
&format!("{}/token", server.uri()),
);
let env = MapEnv::new().with(DOCS_API_URL, &server.uri());
let docs = DocsClient::from_drive_client_with(&env, &drive).unwrap();
(drive, docs)
}
#[test]
fn parse_suggestions_view_defaults_to_current_access() {
assert_eq!(
parse_suggestions_view(None).unwrap(),
SuggestionsViewMode::DefaultForCurrentAccess
);
assert_eq!(
parse_suggestions_view(Some("default")).unwrap(),
SuggestionsViewMode::DefaultForCurrentAccess
);
}
#[test]
fn parse_suggestions_view_accepts_known_values_case_insensitively() {
assert_eq!(
parse_suggestions_view(Some("Inline")).unwrap(),
SuggestionsViewMode::Inline
);
assert_eq!(
parse_suggestions_view(Some("ACCEPTED")).unwrap(),
SuggestionsViewMode::PreviewAccepted
);
assert_eq!(
parse_suggestions_view(Some("without")).unwrap(),
SuggestionsViewMode::PreviewWithoutSuggestions
);
}
#[test]
fn parse_suggestions_view_rejects_unknown_values() {
let err = parse_suggestions_view(Some("bogus")).unwrap_err();
assert!(err.to_string().contains("unknown suggestions_view"));
}
#[tokio::test]
async fn run_docs_info_returns_title_revision_and_outline_as_yaml() {
let server = MockServer::start().await;
let (drive, docs) = docs_and_drive_clients(&server).await;
Mock::given(method("GET"))
.and(path("/v1/documents/d1"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"documentId": "d1",
"title": "Design Doc",
"revisionId": "rev-abc",
"body": {"content": []},
})))
.expect(1)
.mount(&server)
.await;
let params = DriveDocsInfoParams {
document_id: "d1".to_string(),
account: None,
};
let yaml = run_docs_info(&drive, &docs, ¶ms).await.unwrap();
assert!(yaml.contains("Design Doc"), "{yaml}");
assert!(yaml.contains("rev-abc"), "{yaml}");
}
#[tokio::test]
async fn run_docs_info_explains_a_non_doc_target() {
let server = MockServer::start().await;
let (drive, docs) = docs_and_drive_clients(&server).await;
Mock::given(method("GET"))
.and(path("/v1/documents/s1"))
.respond_with(ResponseTemplate::new(404).set_body_json(serde_json::json!({
"error": {"code": 404, "message": "Requested entity was not found.",
"status": "NOT_FOUND"},
})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/drive/v3/files/s1"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"id": "s1",
"name": "Budget",
"mimeType": "application/vnd.google-apps.spreadsheet",
})))
.mount(&server)
.await;
let params = DriveDocsInfoParams {
document_id: "s1".to_string(),
account: None,
};
let err = run_docs_info(&drive, &docs, ¶ms).await.unwrap_err();
let text = err.to_string();
assert!(text.contains("not a Google Doc"), "{text}");
assert!(text.contains("drive sheets read"), "{text}");
}
fn read_params(tab: Option<&str>) -> DriveDocsReadParams {
DriveDocsReadParams {
document_id: "d1".to_string(),
tab: tab.map(str::to_string),
suggestions_view: None,
output_file: None,
account: None,
}
}
fn document_body() -> serde_json::Value {
serde_json::json!({
"documentId": "d1",
"title": "Design Doc",
"revisionId": "rev-1",
"body": {"content": [
{
"startIndex": 1, "endIndex": 10,
"paragraph": {
"elements": [{"textRun": {"content": "Overview\n"}}],
"paragraphStyle": {"namedStyleType": "HEADING_1"},
},
},
]},
})
}
#[tokio::test]
async fn run_docs_read_returns_elements_with_revision_id() {
let server = MockServer::start().await;
let (drive, docs) = docs_and_drive_clients(&server).await;
Mock::given(method("GET"))
.and(path("/v1/documents/d1"))
.respond_with(ResponseTemplate::new(200).set_body_json(document_body()))
.mount(&server)
.await;
let yaml = run_docs_read(&drive, &docs, &read_params(None))
.await
.unwrap();
assert!(yaml.contains("rev-1"), "{yaml}");
assert!(yaml.contains("Overview"), "{yaml}");
}
#[tokio::test]
async fn run_docs_read_sends_the_requested_suggestions_view_mode() {
let server = MockServer::start().await;
let (drive, docs) = docs_and_drive_clients(&server).await;
Mock::given(method("GET"))
.and(path("/v1/documents/d1"))
.and(query_param(
"suggestionsViewMode",
"PREVIEW_SUGGESTIONS_ACCEPTED",
))
.respond_with(ResponseTemplate::new(200).set_body_json(document_body()))
.expect(1)
.mount(&server)
.await;
let mut params = read_params(None);
params.suggestions_view = Some("accepted".to_string());
run_docs_read(&drive, &docs, ¶ms).await.unwrap();
}
#[tokio::test]
async fn run_docs_read_restricts_to_one_tab() {
let server = MockServer::start().await;
let (drive, docs) = docs_and_drive_clients(&server).await;
Mock::given(method("GET"))
.and(path("/v1/documents/d1"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"documentId": "d1",
"revisionId": "rev-1",
"tabs": [
{
"tabProperties": {"tabId": "t.0", "title": "One"},
"documentTab": {"body": {"content": [
{"startIndex": 1, "endIndex": 5,
"paragraph": {"elements": [{"textRun": {"content": "a\n"}}]}},
]}},
},
{
"tabProperties": {"tabId": "t.1", "title": "Two"},
"documentTab": {"body": {"content": [
{"startIndex": 1, "endIndex": 5,
"paragraph": {"elements": [{"textRun": {"content": "b\n"}}]}},
]}},
},
],
})))
.mount(&server)
.await;
let yaml = run_docs_read(&drive, &docs, &read_params(Some("t.1")))
.await
.unwrap();
assert!(yaml.contains("t.1"), "{yaml}");
assert!(!yaml.contains("t.0"), "{yaml}");
}
#[tokio::test]
async fn run_docs_read_rejects_an_unknown_tab_naming_the_real_ones() {
let server = MockServer::start().await;
let (drive, docs) = docs_and_drive_clients(&server).await;
Mock::given(method("GET"))
.and(path("/v1/documents/d1"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"documentId": "d1",
"tabs": [
{
"tabProperties": {"tabId": "t.0", "title": "One"},
"documentTab": {"body": {"content": []}},
},
],
})))
.mount(&server)
.await;
let err = run_docs_read(&drive, &docs, &read_params(Some("missing")))
.await
.unwrap_err();
assert!(err.to_string().contains("t.0"), "{err}");
}
#[tokio::test]
async fn run_docs_read_rejects_an_unknown_suggestions_view_before_any_request() {
let server = MockServer::start().await;
let (drive, docs) = docs_and_drive_clients(&server).await;
let mut params = read_params(None);
params.suggestions_view = Some("bogus".to_string());
let err = run_docs_read(&drive, &docs, ¶ms).await.unwrap_err();
assert!(
err.to_string().contains("unknown suggestions_view"),
"{err}"
);
}
#[tokio::test]
async fn run_docs_read_writes_output_file_and_returns_summary() {
let server = MockServer::start().await;
let (drive, docs) = docs_and_drive_clients(&server).await;
Mock::given(method("GET"))
.and(path("/v1/documents/d1"))
.respond_with(ResponseTemplate::new(200).set_body_json(document_body()))
.mount(&server)
.await;
let tmp = tempfile::tempdir().unwrap();
let out_path = tmp.path().join("out.yaml");
let mut params = read_params(None);
params.output_file = Some(out_path.to_str().unwrap().to_string());
let summary = run_docs_read(&drive, &docs, ¶ms).await.unwrap();
assert!(!summary.contains("Overview"), "{summary}");
assert!(summary.contains("format: yaml"), "{summary}");
let written = std::fs::read_to_string(&out_path).unwrap();
assert!(written.contains("Overview"), "{written}");
}
#[tokio::test]
async fn run_docs_read_keeps_the_original_error_when_classification_fails() {
let server = MockServer::start().await;
let (drive, docs) = docs_and_drive_clients(&server).await;
Mock::given(method("GET"))
.and(path("/v1/documents/d1"))
.respond_with(ResponseTemplate::new(500).set_body_json(serde_json::json!({
"error": {"code": 500, "message": "backend error", "status": "INTERNAL"},
})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/drive/v3/files/d1"))
.respond_with(ResponseTemplate::new(500))
.mount(&server)
.await;
let err = run_docs_read(&drive, &docs, &read_params(None))
.await
.unwrap_err();
assert!(err.to_string().contains("backend error"), "{err}");
}
#[tokio::test(flavor = "current_thread")]
async fn drive_docs_info_handler_propagates_credentials_error() {
let guard = EnvGuard::take();
let _dir = guard.clear_credentials();
let server = OmniDevServer::new();
let err = server
.drive_docs_info(Parameters(DriveDocsInfoParams {
document_id: "d1".to_string(),
account: None,
}))
.await
.unwrap_err();
assert!(err.message.contains("not configured"));
}
#[tokio::test(flavor = "current_thread")]
async fn drive_docs_read_handler_propagates_credentials_error() {
let guard = EnvGuard::take();
let _dir = guard.clear_credentials();
let server = OmniDevServer::new();
let err = server
.drive_docs_read(Parameters(read_params(None)))
.await
.unwrap_err();
assert!(err.message.contains("not configured"));
}
}