yenru0/zed
1
0 Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
fb38178d74c871dbfa60e80038ca5c5d646b1cb7
zed/crates/agent_ui/src/buffer_codegen.rs
T
morgankreyKirill BulatovChristopher Biscardi
 28e52a2879 Send session headers on all OpenCode requests (#63702) 
Fixes https://github.com/zed-industries/zed/issues/63672

OpenCode Go now requires the `x-opencode-session` header on every model
request. Zed previously sent the header only when a request already had
a thread ID, so inline assistance, terminal assistance, and other
standalone requests could omit it and risk rejection.

This change makes the OpenCode provider add the header to every request.
Existing thread and assist session IDs are reused when available.
Standalone requests receive a generated non-empty ID. Thread title and
summary requests now carry the thread ID so related requests keep the
same session scope.

The change also adds tests for missing, empty, and invalid IDs, plus a
wire-level test that verifies the header reaches the HTTP client. The
provider test file stays separate so the implementation remains under
1,000 lines.

Testing:

- `cargo fmt --all -- --check`
- Focused OpenCode, thread summary, and thread title tests
- `cargo check -p agent_ui -p sidebar`
- `GITHUB_ACTIONS=1 ./script/clippy -p language_models`
- `git diff --check`

Release Notes:

- Fixed OpenCode Go requests that could fail when Zed omitted the
session header.

---------

Co-authored-by: Kirill Bulatov <kirill@zed.dev>
Co-authored-by: Christopher Biscardi <chris@christopherbiscardi.com>
2026-09-03 18:29:10 +00:00

2675 lines
101 KiB
Rust
Raw Blame History

use crate::{context::LoadedContext, inline_prompt_editor::CodegenStatus};
use agent_settings::AgentSettings;
use anyhow::{Context as _, Result, anyhow};
use collections::HashSet;
use editor::{Anchor, AnchorRangeExt, MultiBuffer, MultiBufferSnapshot, ToOffset as _, ToPoint};
use futures::FutureExt;
use futures::{
SinkExt, Stream, StreamExt, TryStreamExt as _,
channel::mpsc,
future::{LocalBoxFuture, Shared},
join,
stream::BoxStream,
};
use gpui::{App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, Subscription, Task};
use language::{
Buffer, BufferEditSource, IndentKind, LanguageName, Point, TransactionId, line_diff,
};
use language_model::{
CompletionIntent, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent,
LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
LanguageModelRequestTool, LanguageModelTextStream, LanguageModelToolChoice,
LanguageModelToolUse, LanguageModelToolUseId, Role, StopReason, TokenUsage,
};
use language_models::provider::anthropic::telemetry::{
AnthropicCompletionType, AnthropicEventData, AnthropicEventReporter, AnthropicEventType,
};
use multi_buffer::MultiBufferRow;
use parking_lot::Mutex;
use prompt_store::PromptBuilder;
use rope::Rope;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use settings::Settings as _;
use std::{
cmp,
future::Future,
iter,
ops::{Range, RangeInclusive},
pin::Pin,
sync::Arc,
task::{self, Poll},
time::Instant,
};
use streaming_diff::{CharOperation, LineDiff, LineOperation, StreamingDiff};
use uuid::Uuid;
/// Use this tool when you cannot or should not make a rewrite. This includes:
/// - The user's request is unclear, ambiguous, or nonsensical
/// - The requested change cannot be made by only editing the <rewrite_this> section
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
pub struct FailureMessageInput {
/// A brief message to the user explaining why you're unable to fulfill the request or to ask a question about the request.
#[serde(default)]
pub message: String,
}
/// Replaces text in <rewrite_this></rewrite_this> tags with your replacement_text.
/// Only use this tool when you are confident you understand the user's request and can fulfill it
/// by editing the marked section.
#[derive(Debug, Serialize, Deserialize, JsonSchema)]
pub struct RewriteSectionInput {
/// The text to replace the section with.
#[serde(default)]
pub replacement_text: String,
}
pub struct BufferCodegen {
alternatives: Vec<Entity<CodegenAlternative>>,
pub active_alternative: usize,
seen_alternatives: HashSet<usize>,
subscriptions: Vec<Subscription>,
buffer: Entity<MultiBuffer>,
range: Range<Anchor>,
initial_transaction_id: Option<TransactionId>,
builder: Arc<PromptBuilder>,
pub is_insertion: bool,
session_id: Uuid,
}
pub const REWRITE_SECTION_TOOL_NAME: &str = "rewrite_section";
pub const FAILURE_MESSAGE_TOOL_NAME: &str = "failure_message";
impl BufferCodegen {
pub fn new(
buffer: Entity<MultiBuffer>,
range: Range<Anchor>,
initial_transaction_id: Option<TransactionId>,
session_id: Uuid,
builder: Arc<PromptBuilder>,
cx: &mut Context<Self>,
) -> Self {
let codegen = cx.new(|cx| {
CodegenAlternative::new(
buffer.clone(),
range.clone(),
false,
builder.clone(),
session_id,
cx,
)
});
let mut this = Self {
is_insertion: range.to_offset(&buffer.read(cx).snapshot(cx)).is_empty(),
alternatives: vec![codegen],
active_alternative: 0,
seen_alternatives: HashSet::default(),
subscriptions: Vec::new(),
buffer,
range,
initial_transaction_id,
builder,
session_id,
};
this.activate(0, cx);
this
}
fn subscribe_to_alternative(&mut self, cx: &mut Context<Self>) {
let codegen = self.active_alternative().clone();
self.subscriptions.clear();
self.subscriptions
.push(cx.observe(&codegen, |_, _, cx| cx.notify()));
self.subscriptions
.push(cx.subscribe(&codegen, |_, _, event, cx| cx.emit(*event)));
}
pub fn active_completion(&self, cx: &App) -> Option<String> {
self.active_alternative().read(cx).current_completion()
}
pub fn active_alternative(&self) -> &Entity<CodegenAlternative> {
&self.alternatives[self.active_alternative]
}
pub fn language_name(&self, cx: &App) -> Option<LanguageName> {
self.active_alternative().read(cx).language_name(cx)
}
pub fn status<'a>(&self, cx: &'a App) -> &'a CodegenStatus {
&self.active_alternative().read(cx).status
}
pub fn alternative_count(&self, cx: &App) -> usize {
LanguageModelRegistry::read_global(cx)
.inline_alternative_models()
.len()
+ 1
}
pub fn cycle_prev(&mut self, cx: &mut Context<Self>) {
let next_active_ix = if self.active_alternative == 0 {
self.alternatives.len() - 1
} else {
self.active_alternative - 1
};
self.activate(next_active_ix, cx);
}
pub fn cycle_next(&mut self, cx: &mut Context<Self>) {
let next_active_ix = (self.active_alternative + 1) % self.alternatives.len();
self.activate(next_active_ix, cx);
}
fn activate(&mut self, index: usize, cx: &mut Context<Self>) {
self.active_alternative()
.update(cx, |codegen, cx| codegen.set_active(false, cx));
self.seen_alternatives.insert(index);
self.active_alternative = index;
self.active_alternative()
.update(cx, |codegen, cx| codegen.set_active(true, cx));
self.subscribe_to_alternative(cx);
cx.notify();
}
pub fn start(
&mut self,
primary_model: Arc<dyn LanguageModel>,
user_prompt: String,
context_task: Shared<Task<Option<LoadedContext>>>,
cx: &mut Context<Self>,
) -> Result<()> {
let alternative_models = LanguageModelRegistry::read_global(cx)
.inline_alternative_models()
.to_vec();
self.active_alternative()
.update(cx, |alternative, cx| alternative.undo(cx));
self.activate(0, cx);
self.alternatives.truncate(1);
for _ in 0..alternative_models.len() {
self.alternatives.push(cx.new(|cx| {
CodegenAlternative::new(
self.buffer.clone(),
self.range.clone(),
false,
self.builder.clone(),
self.session_id,
cx,
)
}));
}
for (model, alternative) in iter::once(primary_model)
.chain(alternative_models)
.zip(&self.alternatives)
{
alternative.update(cx, |alternative, cx| {
alternative.start(user_prompt.clone(), context_task.clone(), model.clone(), cx)
})?;
}
Ok(())
}
pub fn stop(&mut self, cx: &mut Context<Self>) {
for codegen in &self.alternatives {
codegen.update(cx, |codegen, cx| codegen.stop(cx));
}
}
pub fn undo(&mut self, cx: &mut Context<Self>) {
self.active_alternative()
.update(cx, |codegen, cx| codegen.undo(cx));
self.buffer.update(cx, |buffer, cx| {
if let Some(transaction_id) = self.initial_transaction_id.take() {
buffer.undo_transaction(transaction_id, cx);
buffer.refresh_preview(cx);
}
});
}
pub fn buffer(&self, cx: &App) -> Entity<MultiBuffer> {
self.active_alternative().read(cx).buffer.clone()
}
pub fn old_buffer(&self, cx: &App) -> Entity<Buffer> {
self.active_alternative().read(cx).old_buffer.clone()
}
pub fn snapshot(&self, cx: &App) -> MultiBufferSnapshot {
self.active_alternative().read(cx).snapshot.clone()
}
pub fn edit_position(&self, cx: &App) -> Option<Anchor> {
self.active_alternative().read(cx).edit_position
}
pub fn diff<'a>(&self, cx: &'a App) -> &'a Diff {
&self.active_alternative().read(cx).diff
}
pub fn last_equal_ranges<'a>(&self, cx: &'a App) -> &'a [Range<Anchor>] {
self.active_alternative().read(cx).last_equal_ranges()
}
pub fn selected_text<'a>(&self, cx: &'a App) -> Option<&'a str> {
self.active_alternative().read(cx).selected_text()
}
pub fn session_id(&self) -> Uuid {
self.session_id
}
}
impl EventEmitter<CodegenEvent> for BufferCodegen {}
pub struct CodegenAlternative {
buffer: Entity<MultiBuffer>,
old_buffer: Entity<Buffer>,
snapshot: MultiBufferSnapshot,
edit_position: Option<Anchor>,
range: Range<Anchor>,
last_equal_ranges: Vec<Range<Anchor>>,
transformation_transaction_id: Option<TransactionId>,
status: CodegenStatus,
generation: Task<()>,
diff: Diff,
_subscription: gpui::Subscription,
builder: Arc<PromptBuilder>,
active: bool,
edits: Vec<(Range<Anchor>, String)>,
line_operations: Vec<LineOperation>,
elapsed_time: Option<f64>,
completion: Option<String>,
selected_text: Option<String>,
pub message_id: Option<String>,
session_id: Uuid,
pub description: Option<String>,
pub failure: Option<String>,
}
impl EventEmitter<CodegenEvent> for CodegenAlternative {}
impl CodegenAlternative {
pub fn new(
buffer: Entity<MultiBuffer>,
range: Range<Anchor>,
active: bool,
builder: Arc<PromptBuilder>,
session_id: Uuid,
cx: &mut Context<Self>,
) -> Self {
let snapshot = buffer.read(cx).snapshot(cx);
let (old_buffer, _, _) = snapshot
.range_to_buffer_ranges(range.start..range.end)
.pop()
.unwrap();
let old_buffer = cx.new(|cx| {
let text = old_buffer.as_rope().clone();
let line_ending = old_buffer.line_ending();
let language = old_buffer.language().cloned();
let language_registry = buffer
.read(cx)
.buffer(old_buffer.remote_id())
.unwrap()
.read(cx)
.language_registry();
let mut buffer = Buffer::local_normalized(text, line_ending, cx);
buffer.set_language(language, cx);
if let Some(language_registry) = language_registry {
buffer.set_language_registry(language_registry);
}
buffer
});
Self {
buffer: buffer.clone(),
old_buffer,
edit_position: None,
message_id: None,
snapshot,
last_equal_ranges: Default::default(),
transformation_transaction_id: None,
status: CodegenStatus::Idle,
generation: Task::ready(()),
diff: Diff::default(),
builder,
active: active,
edits: Vec::new(),
line_operations: Vec::new(),
range,
elapsed_time: None,
completion: None,
selected_text: None,
session_id,
description: None,
failure: None,
_subscription: cx.subscribe(&buffer, Self::handle_buffer_event),
}
}
pub fn language_name(&self, cx: &App) -> Option<LanguageName> {
self.old_buffer
.read(cx)
.language()
.map(|language| language.name())
}
pub fn set_active(&mut self, active: bool, cx: &mut Context<Self>) {
if active != self.active {
self.active = active;
if self.active {
let edits = self.edits.clone();
self.apply_edits(edits, cx);
if matches!(self.status, CodegenStatus::Pending) {
let line_operations = self.line_operations.clone();
self.reapply_line_based_diff(line_operations, cx);
} else {
self.reapply_batch_diff(cx).detach();
}
} else if let Some(transaction_id) = self.transformation_transaction_id.take() {
self.buffer.update(cx, |buffer, cx| {
buffer.undo_transaction(transaction_id, cx);
buffer.forget_transaction(transaction_id, cx);
});
}
}
}
fn handle_buffer_event(
&mut self,
_buffer: Entity<MultiBuffer>,
event: &multi_buffer::Event,
cx: &mut Context<Self>,
) {
if let multi_buffer::Event::TransactionUndone { transaction_id } = event
&& self.transformation_transaction_id == Some(*transaction_id)
{
self.transformation_transaction_id = None;
self.generation = Task::ready(());
cx.emit(CodegenEvent::Undone);
}
}
pub fn last_equal_ranges(&self) -> &[Range<Anchor>] {
&self.last_equal_ranges
}
pub fn use_streaming_tools(model: &dyn LanguageModel, cx: &App) -> bool {
model.supports_streaming_tools()
&& AgentSettings::get_global(cx).inline_assistant_use_streaming_tools
}
pub fn start(
&mut self,
user_prompt: String,
context_task: Shared<Task<Option<LoadedContext>>>,
model: Arc<dyn LanguageModel>,
cx: &mut Context<Self>,
) -> Result<()> {
// Clear the model explanation since the user has started a new generation.
self.description = None;
self.failure = None;
if let Some(transformation_transaction_id) = self.transformation_transaction_id.take() {
self.buffer.update(cx, |buffer, cx| {
buffer.undo_transaction(transformation_transaction_id, cx);
});
}
self.edit_position = Some(self.range.start.bias_right(&self.snapshot));
if Self::use_streaming_tools(model.as_ref(), cx) {
let request = self.build_request(&model, user_prompt, context_task, cx)?;
let completion_events = cx.spawn({
let model = model.clone();
async move |_, cx| model.stream_completion(request.await, cx).await
});
self.generation = self.handle_completion(model, completion_events, cx);
} else {
let stream: LocalBoxFuture<Result<LanguageModelTextStream>> =
if user_prompt.trim().to_lowercase() == "delete" {
async { Ok(LanguageModelTextStream::default()) }.boxed_local()
} else {
let request = self.build_request(&model, user_prompt, context_task, cx)?;
cx.spawn({
let model = model.clone();
async move |_, cx| {
Ok(model.stream_completion_text(request.await, cx).await?)
}
})
.boxed_local()
};
self.generation =
self.handle_stream(model, /* strip_invalid_spans: */ true, stream, cx);
}
Ok(())
}
fn build_request_tools(
&self,
model: &Arc<dyn LanguageModel>,
user_prompt: String,
context_task: Shared<Task<Option<LoadedContext>>>,
cx: &mut App,
) -> Result<Task<LanguageModelRequest>> {
let buffer = self.buffer.read(cx).snapshot(cx);
let language = buffer.language_at(self.range.start);
let language_name = if let Some(language) = language.as_ref() {
if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
None
} else {
Some(language.name())
}
} else {
None
};
let language_name = language_name.as_ref();
let start = buffer.point_to_buffer_offset(self.range.start);
let end = buffer.point_to_buffer_offset(self.range.end);
let (buffer, range) = if let Some((start, end)) = start.zip(end) {
let (start_buffer, start_buffer_offset) = start;
let (end_buffer, end_buffer_offset) = end;
if start_buffer.remote_id() == end_buffer.remote_id() {
(start_buffer.clone(), start_buffer_offset..end_buffer_offset)
} else {
anyhow::bail!("invalid transformation range");
}
} else {
anyhow::bail!("invalid transformation range");
};
let system_prompt = self
.builder
.generate_inline_transformation_prompt_tools(
language_name,
buffer,
range.start.0..range.end.0,
)
.context("generating content prompt")?;
let temperature = AgentSettings::temperature_for_model(model, cx);
let session_id = self.session_id.to_string();
let tool_choice = model
.supports_tool_choice(LanguageModelToolChoice::Any)
.then_some(LanguageModelToolChoice::Any);
Ok(cx.spawn(async move |_cx| {
let mut messages = vec![LanguageModelRequestMessage {
role: Role::System,
content: vec![system_prompt.into()],
cache: false,
reasoning_details: None,
}];
let mut user_message = LanguageModelRequestMessage {
role: Role::User,
content: Vec::new(),
cache: false,
reasoning_details: None,
};
if let Some(context) = context_task.await {
context.add_to_request_message(&mut user_message);
}
user_message.content.push(user_prompt.into());
messages.push(user_message);
let tool_schema = |mut schema: serde_json::Value| {
language_model::tool_schema::normalize_tool_schema(&mut schema);
schema
};
let tools = vec![
LanguageModelRequestTool::function(
REWRITE_SECTION_TOOL_NAME.to_string(),
"Replaces text in <rewrite_this></rewrite_this> tags with your replacement_text.".to_string(),
tool_schema(language_model::tool_schema::root_schema_for::<RewriteSectionInput>().to_value()),
false,
),
LanguageModelRequestTool::function(
FAILURE_MESSAGE_TOOL_NAME.to_string(),
"Use this tool to provide a message to the user when you're unable to complete a task.".to_string(),
tool_schema(language_model::tool_schema::root_schema_for::<FailureMessageInput>().to_value()),
false,
),
];
LanguageModelRequest {
thread_id: Some(session_id),
prompt_id: None,
intent: Some(CompletionIntent::InlineAssist),
tools,
tool_choice,
stop: Vec::new(),
temperature,
messages,
thinking_allowed: false,
thinking_effort: None,
speed: None,
compact_at_tokens: None,
}
}))
}
fn build_request(
&self,
model: &Arc<dyn LanguageModel>,
user_prompt: String,
context_task: Shared<Task<Option<LoadedContext>>>,
cx: &mut App,
) -> Result<Task<LanguageModelRequest>> {
if Self::use_streaming_tools(model.as_ref(), cx) {
return self.build_request_tools(model, user_prompt, context_task, cx);
}
let buffer = self.buffer.read(cx).snapshot(cx);
let language = buffer.language_at(self.range.start);
let language_name = if let Some(language) = language.as_ref() {
if Arc::ptr_eq(language, &language::PLAIN_TEXT) {
None
} else {
Some(language.name())
}
} else {
None
};
let language_name = language_name.as_ref();
let start = buffer.point_to_buffer_offset(self.range.start);
let end = buffer.point_to_buffer_offset(self.range.end);
let (buffer, range) = if let Some((start, end)) = start.zip(end) {
let (start_buffer, start_buffer_offset) = start;
let (end_buffer, end_buffer_offset) = end;
if start_buffer.remote_id() == end_buffer.remote_id() {
(start_buffer.clone(), start_buffer_offset..end_buffer_offset)
} else {
anyhow::bail!("invalid transformation range");
}
} else {
anyhow::bail!("invalid transformation range");
};
let prompt = self
.builder
.generate_inline_transformation_prompt(
user_prompt,
language_name,
buffer,
range.start.0..range.end.0,
)
.context("generating content prompt")?;
let temperature = AgentSettings::temperature_for_model(model, cx);
let session_id = self.session_id.to_string();
Ok(cx.spawn(async move |_cx| {
let mut request_message = LanguageModelRequestMessage {
role: Role::User,
content: Vec::new(),
cache: false,
reasoning_details: None,
};
if let Some(context) = context_task.await {
context.add_to_request_message(&mut request_message);
}
request_message.content.push(prompt.into());
LanguageModelRequest {
thread_id: Some(session_id),
prompt_id: None,
intent: Some(CompletionIntent::InlineAssist),
tools: Vec::new(),
tool_choice: None,
stop: Vec::new(),
temperature,
messages: vec![request_message],
thinking_allowed: false,
thinking_effort: None,
speed: None,
compact_at_tokens: None,
}
}))
}
pub fn handle_stream(
&mut self,
model: Arc<dyn LanguageModel>,
strip_invalid_spans: bool,
stream: impl 'static + Future<Output = Result<LanguageModelTextStream>>,
cx: &mut Context<Self>,
) -> Task<()> {
let anthropic_reporter = AnthropicEventReporter::new(&model, cx);
let session_id = self.session_id;
let model_telemetry_id = model.telemetry_id();
let model_provider_id = model.provider_id().to_string();
let start_time = Instant::now();
// Make a new snapshot and re-resolve anchor in case the document was modified.
// This can happen often if the editor loses focus and is saved + reformatted,
// as in https://github.com/zed-industries/zed/issues/39088
self.snapshot = self.buffer.read(cx).snapshot(cx);
self.range = self.snapshot.anchor_after(self.range.start)
..self.snapshot.anchor_after(self.range.end);
let snapshot = self.snapshot.clone();
let selected_text = snapshot
.text_for_range(self.range.start..self.range.end)
.collect::<Rope>();
self.selected_text = Some(selected_text.to_string());
let selection_start = self.range.start.to_point(&snapshot);
// Start with the indentation of the first line in the selection
let mut suggested_line_indent = snapshot
.suggested_indents(selection_start.row..=selection_start.row, cx)
.into_values()
.next()
.unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
// If the first line in the selection does not have indentation, check the following lines
if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {
for row in selection_start.row..=self.range.end.to_point(&snapshot).row {
let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));
// Prefer tabs if a line in the selection uses tabs as indentation
if line_indent.kind == IndentKind::Tab {
suggested_line_indent.kind = IndentKind::Tab;
break;
}
}
}
let language_name = {
let multibuffer = self.buffer.read(cx);
let snapshot = multibuffer.snapshot(cx);
let ranges = snapshot.range_to_buffer_ranges(self.range.start..self.range.end);
ranges
.first()
.and_then(|(buffer, _, _)| buffer.language())
.map(|language| language.name())
};
self.diff = Diff::default();
self.status = CodegenStatus::Pending;
let mut edit_start = self.range.start.to_offset(&snapshot);
let completion = Arc::new(Mutex::new(String::new()));
let completion_clone = completion.clone();
cx.notify();
cx.spawn(async move |codegen, cx| {
let stream = stream.await;
let token_usage = stream
.as_ref()
.ok()
.map(|stream| stream.last_token_usage.clone());
let message_id = stream
.as_ref()
.ok()
.and_then(|stream| stream.message_id.clone());
let generate = async {
let model_telemetry_id = model_telemetry_id.clone();
let model_provider_id = model_provider_id.clone();
let (mut diff_tx, mut diff_rx) = mpsc::channel(1);
let message_id = message_id.clone();
let line_based_stream_diff: Task<anyhow::Result<()>> = cx.background_spawn({
let anthropic_reporter = anthropic_reporter.clone();
let language_name = language_name.clone();
async move {
let mut response_latency = None;
let request_start = Instant::now();
let diff = async {
let raw_stream = stream?.stream.map_err(|error| error.into());
let stripped;
let mut chunks: Pin<Box<dyn Stream<Item = Result<String>> + Send>> =
if strip_invalid_spans {
stripped = StripInvalidSpans::new(raw_stream);
Box::pin(stripped)
} else {
Box::pin(raw_stream)
};
let mut diff = StreamingDiff::new(selected_text.to_string());
let mut line_diff = LineDiff::default();
let mut new_text = String::new();
let mut base_indent = None;
let mut line_indent = None;
let mut first_line = true;
while let Some(chunk) = chunks.next().await {
if response_latency.is_none() {
response_latency = Some(request_start.elapsed());
}
let chunk = chunk?;
completion_clone.lock().push_str(&chunk);
let mut lines = chunk.split('\n').peekable();
while let Some(line) = lines.next() {
new_text.push_str(line);
if line_indent.is_none()
&& let Some(non_whitespace_ch_ix) =
new_text.find(|ch: char| !ch.is_whitespace())
{
line_indent = Some(non_whitespace_ch_ix);
base_indent = base_indent.or(line_indent);
let line_indent = line_indent.unwrap();
let base_indent = base_indent.unwrap();
let indent_delta = line_indent as i32 - base_indent as i32;
let mut corrected_indent_len = cmp::max(
0,
suggested_line_indent.len as i32 + indent_delta,
)
as usize;
if first_line {
corrected_indent_len = corrected_indent_len
.saturating_sub(selection_start.column as usize);
}
let indent_char = suggested_line_indent.char();
let mut indent_buffer = [0; 4];
let indent_str =
indent_char.encode_utf8(&mut indent_buffer);
new_text.replace_range(
..line_indent,
&indent_str.repeat(corrected_indent_len),
);
}
if line_indent.is_some() {
let char_ops = diff.push_new(&new_text);
line_diff.push_char_operations(&char_ops, &selected_text);
diff_tx
.send((char_ops, line_diff.line_operations()))
.await?;
new_text.clear();
}
if lines.peek().is_some() {
let char_ops = diff.push_new("\n");
line_diff.push_char_operations(&char_ops, &selected_text);
diff_tx
.send((char_ops, line_diff.line_operations()))
.await?;
if line_indent.is_none() {
// Don't write out the leading indentation in empty lines on the next line
// This is the case where the above if statement didn't clear the buffer
new_text.clear();
}
line_indent = None;
first_line = false;
}
}
}
let mut char_ops = diff.push_new(&new_text);
char_ops.extend(diff.finish());
line_diff.push_char_operations(&char_ops, &selected_text);
line_diff.finish(&selected_text);
diff_tx
.send((char_ops, line_diff.line_operations()))
.await?;
anyhow::Ok(())
};
let result = diff.await;
let error_message = result.as_ref().err().map(|error| error.to_string());
telemetry::event!(
"Assistant Responded",
kind = "inline",
phase = "response",
session_id = session_id.to_string(),
model = model_telemetry_id,
model_provider = model_provider_id,
language_name = language_name.as_ref().map(|n| n.to_string()),
message_id = message_id.as_deref(),
response_latency = response_latency,
error_message = error_message.as_deref(),
);
anthropic_reporter.report(AnthropicEventData {
completion_type: AnthropicCompletionType::Editor,
event: AnthropicEventType::Response,
language_name: language_name.map(|n| n.to_string()),
message_id,
});
result?;
Ok(())
}
});
while let Some((char_ops, line_ops)) = diff_rx.next().await {
codegen.update(cx, |codegen, cx| {
codegen.last_equal_ranges.clear();
let edits = char_ops
.into_iter()
.filter_map(|operation| match operation {
CharOperation::Insert { text } => {
let edit_start = snapshot.anchor_after(edit_start);
Some((edit_start..edit_start, text))
}
CharOperation::Delete { bytes } => {
let edit_end = edit_start + bytes;
let edit_range = snapshot.anchor_after(edit_start)
..snapshot.anchor_before(edit_end);
edit_start = edit_end;
Some((edit_range, String::new()))
}
CharOperation::Keep { bytes } => {
let edit_end = edit_start + bytes;
let edit_range = snapshot.anchor_after(edit_start)
..snapshot.anchor_before(edit_end);
edit_start = edit_end;
codegen.last_equal_ranges.push(edit_range);
None
}
})
.collect::<Vec<_>>();
if codegen.active {
codegen.apply_edits(edits.iter().cloned(), cx);
codegen.reapply_line_based_diff(line_ops.iter().cloned(), cx);
}
codegen.edits.extend(edits);
codegen.line_operations = line_ops;
codegen.edit_position = Some(snapshot.anchor_after(edit_start));
cx.notify();
})?;
}
// Streaming stopped and we have the new text in the buffer, and a line-based diff applied for the whole new buffer.
// That diff is not what a regular diff is and might look unexpected, ergo apply a regular diff.
// It's fine to apply even if the rest of the line diffing fails, as no more hunks are coming through `diff_rx`.
let batch_diff_task =
codegen.update(cx, |codegen, cx| codegen.reapply_batch_diff(cx))?;
let (line_based_stream_diff, ()) = join!(line_based_stream_diff, batch_diff_task);
line_based_stream_diff?;
anyhow::Ok(())
};
let result = generate.await;
let elapsed_time = start_time.elapsed().as_secs_f64();
codegen
.update(cx, |this, cx| {
this.message_id = message_id;
this.last_equal_ranges.clear();
if let Err(error) = result {
this.status = CodegenStatus::Error(error);
} else {
this.status = CodegenStatus::Done;
}
this.elapsed_time = Some(elapsed_time);
this.completion = Some(completion.lock().clone());
if let Some(usage) = token_usage {
let usage = usage.lock();
telemetry::event!(
"Inline Assistant Completion",
model = model_telemetry_id,
model_provider = model_provider_id,
input_tokens = usage.input_tokens,
output_tokens = usage.output_tokens,
)
}
cx.emit(CodegenEvent::Finished);
cx.notify();
})
.ok();
})
}
pub fn current_completion(&self) -> Option<String> {
self.completion.clone()
}
#[cfg(any(test, feature = "test-support"))]
pub fn current_description(&self) -> Option<String> {
self.description.clone()
}
#[cfg(any(test, feature = "test-support"))]
pub fn current_failure(&self) -> Option<String> {
self.failure.clone()
}
pub fn selected_text(&self) -> Option<&str> {
self.selected_text.as_deref()
}
pub fn stop(&mut self, cx: &mut Context<Self>) {
self.last_equal_ranges.clear();
if self.diff.is_empty() {
self.status = CodegenStatus::Idle;
} else {
self.status = CodegenStatus::Done;
}
self.generation = Task::ready(());
cx.emit(CodegenEvent::Finished);
cx.notify();
}
pub fn undo(&mut self, cx: &mut Context<Self>) {
self.buffer.update(cx, |buffer, cx| {
if let Some(transaction_id) = self.transformation_transaction_id.take() {
buffer.undo_transaction(transaction_id, cx);
buffer.refresh_preview(cx);
}
});
}
fn apply_edits(
&mut self,
edits: impl IntoIterator<Item = (Range<Anchor>, String)>,
cx: &mut Context<CodegenAlternative>,
) {
let transaction = self.buffer.update(cx, |buffer, cx| {
// Avoid grouping agent edits with user edits.
buffer.finalize_last_transaction(cx);
buffer.start_transaction(cx);
buffer.edit(edits, None, cx);
buffer.end_transaction_with_source(BufferEditSource::Agent, cx)
});
if let Some(transaction) = transaction {
if let Some(first_transaction) = self.transformation_transaction_id {
// Group all agent edits into the first transaction.
self.buffer.update(cx, |buffer, cx| {
buffer.merge_transactions(transaction, first_transaction, cx)
});
} else {
self.transformation_transaction_id = Some(transaction);
self.buffer
.update(cx, |buffer, cx| buffer.finalize_last_transaction(cx));
}
}
}
fn reapply_line_based_diff(
&mut self,
line_operations: impl IntoIterator<Item = LineOperation>,
cx: &mut Context<Self>,
) {
let old_snapshot = self.snapshot.clone();
let old_range = self.range.to_point(&old_snapshot);
let new_snapshot = self.buffer.read(cx).snapshot(cx);
let new_range = self.range.to_point(&new_snapshot);
let mut old_row = old_range.start.row;
let mut new_row = new_range.start.row;
self.diff.deleted_row_ranges.clear();
self.diff.inserted_row_ranges.clear();
for operation in line_operations {
match operation {
LineOperation::Keep { lines } => {
old_row += lines;
new_row += lines;
}
LineOperation::Delete { lines } => {
let old_end_row = old_row + lines - 1;
let new_row = new_snapshot.anchor_before(Point::new(new_row, 0));
if let Some((_, last_deleted_row_range)) =
self.diff.deleted_row_ranges.last_mut()
{
if *last_deleted_row_range.end() + 1 == old_row {
*last_deleted_row_range = *last_deleted_row_range.start()..=old_end_row;
} else {
self.diff
.deleted_row_ranges
.push((new_row, old_row..=old_end_row));
}
} else {
self.diff
.deleted_row_ranges
.push((new_row, old_row..=old_end_row));
}
old_row += lines;
}
LineOperation::Insert { lines } => {
let new_end_row = new_row + lines - 1;
let start = new_snapshot.anchor_before(Point::new(new_row, 0));
let end = new_snapshot.anchor_before(Point::new(
new_end_row,
new_snapshot.line_len(MultiBufferRow(new_end_row)),
));
self.diff.inserted_row_ranges.push(start..end);
new_row += lines;
}
}
cx.notify();
}
}
fn reapply_batch_diff(&mut self, cx: &mut Context<Self>) -> Task<()> {
let old_snapshot = self.snapshot.clone();
let old_range = self.range.to_point(&old_snapshot);
let new_snapshot = self.buffer.read(cx).snapshot(cx);
let new_range = self.range.to_point(&new_snapshot);
cx.spawn(async move |codegen, cx| {
let (deleted_row_ranges, inserted_row_ranges) = cx
.background_spawn(async move {
let old_text = old_snapshot
.text_for_range(
Point::new(old_range.start.row, 0)
..Point::new(
old_range.end.row,
old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
),
)
.collect::<String>();
let new_text = new_snapshot
.text_for_range(
Point::new(new_range.start.row, 0)
..Point::new(
new_range.end.row,
new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
),
)
.collect::<String>();
let old_start_row = old_range.start.row;
let new_start_row = new_range.start.row;
let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
let mut inserted_row_ranges = Vec::new();
for (old_rows, new_rows) in line_diff(&old_text, &new_text) {
let old_rows = old_start_row + old_rows.start..old_start_row + old_rows.end;
let new_rows = new_start_row + new_rows.start..new_start_row + new_rows.end;
if !old_rows.is_empty() {
deleted_row_ranges.push((
new_snapshot.anchor_before(Point::new(new_rows.start, 0)),
old_rows.start..=old_rows.end - 1,
));
}
if !new_rows.is_empty() {
let start = new_snapshot.anchor_before(Point::new(new_rows.start, 0));
let new_end_row = new_rows.end - 1;
let end = new_snapshot.anchor_before(Point::new(
new_end_row,
new_snapshot.line_len(MultiBufferRow(new_end_row)),
));
inserted_row_ranges.push(start..end);
}
}
(deleted_row_ranges, inserted_row_ranges)
})
.await;
codegen
.update(cx, |codegen, cx| {
codegen.diff.deleted_row_ranges = deleted_row_ranges;
codegen.diff.inserted_row_ranges = inserted_row_ranges;
cx.notify();
})
.ok();
})
}
fn handle_completion(
&mut self,
model: Arc<dyn LanguageModel>,
completion_stream: Task<
Result<
BoxStream<
'static,
Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
>,
LanguageModelCompletionError,
>,
>,
cx: &mut Context<Self>,
) -> Task<()> {
self.diff = Diff::default();
self.status = CodegenStatus::Pending;
cx.notify();
// Leaving this in generation so that STOP equivalent events are respected even
// while we're still pre-processing the completion event
cx.spawn(async move |codegen, cx| {
let finish_with_status = |status: CodegenStatus, cx: &mut AsyncApp| {
let _ = codegen.update(cx, |this, cx| {
this.status = status;
cx.emit(CodegenEvent::Finished);
cx.notify();
});
};
let mut completion_events = match completion_stream.await {
Ok(events) => events,
Err(err) => {
finish_with_status(CodegenStatus::Error(err.into()), cx);
return;
}
};
enum ToolUseOutput {
Rewrite {
text: String,
},
Failure {
message: String,
is_input_complete: bool,
},
}
struct RewriteState {
tool_use_id: LanguageModelToolUseId,
chars_read: usize,
is_input_complete: bool,
}
let rewrite_state: Arc<Mutex<Option<RewriteState>>> = Arc::new(Mutex::new(None));
let process_rewrite_state = rewrite_state.clone();
let process_tool_use =
move |tool_use: LanguageModelToolUse| -> Result<Option<ToolUseOutput>> {
match tool_use.name.as_ref() {
REWRITE_SECTION_TOOL_NAME => {
let input = match tool_use.input.parse::<RewriteSectionInput>() {
Ok(input) => input,
Err(_) if !tool_use.is_input_complete => return Ok(None),
Err(error) => {
return Err(error.context(format!(
"invalid input for `{REWRITE_SECTION_TOOL_NAME}` tool"
)));
}
};
let mut rewrite_state = process_rewrite_state.lock();
if let Some(state) = rewrite_state.as_ref()
&& state.tool_use_id != tool_use.id
{
return Err(anyhow!(
"Inline assistant called `{REWRITE_SECTION_TOOL_NAME}` \
multiple times"
));
}
let state = rewrite_state.get_or_insert_with(|| RewriteState {
tool_use_id: tool_use.id,
chars_read: 0,
is_input_complete: false,
});
let Some(text_slice) = input.replacement_text.get(state.chars_read..)
else {
return Ok(None);
};
let text = text_slice.to_string();
state.chars_read = input.replacement_text.len();
state.is_input_complete |= tool_use.is_input_complete;
Ok(Some(ToolUseOutput::Rewrite { text }))
}
FAILURE_MESSAGE_TOOL_NAME => {
let mut input = match tool_use.input.parse::<FailureMessageInput>() {
Ok(input) => input,
Err(_) if !tool_use.is_input_complete => return Ok(None),
Err(error) => {
return Err(error.context(format!(
"invalid input for `{FAILURE_MESSAGE_TOOL_NAME}` tool"
)));
}
};
Ok(Some(ToolUseOutput::Failure {
message: std::mem::take(&mut input.message),
is_input_complete: tool_use.is_input_complete,
}))
}
_ => Ok(None),
}
};
let mut message_id = None;
let mut failure_started = false;
let mut failure_complete = false;
let last_token_usage = Arc::new(Mutex::new(TokenUsage::default()));
let first_text = loop {
let Some(event) = completion_events.next().await else {
if failure_complete {
finish_with_status(CodegenStatus::Done, cx);
} else {
finish_with_status(
CodegenStatus::Error(anyhow!(
"Inline assistant response ended without completing \
`{REWRITE_SECTION_TOOL_NAME}` or `{FAILURE_MESSAGE_TOOL_NAME}`"
)),
cx,
);
}
return;
};
match event {
Ok(LanguageModelCompletionEvent::StartMessage { message_id: id }) => {
message_id = Some(id);
}
Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
let output = match process_tool_use(tool_use) {
Ok(output) => output,
Err(error) => {
finish_with_status(CodegenStatus::Error(error), cx);
return;
}
};
if let Some(output) = output {
match output {
ToolUseOutput::Rewrite { text } => {
if failure_started {
finish_with_status(
CodegenStatus::Error(anyhow!(
"Inline assistant called \
`{REWRITE_SECTION_TOOL_NAME}` after \
`{FAILURE_MESSAGE_TOOL_NAME}`"
)),
cx,
);
return;
}
break text;
}
ToolUseOutput::Failure {
message,
is_input_complete,
} => {
failure_started = true;
failure_complete |= is_input_complete;
if is_input_complete
&& codegen
.update(cx, |this, cx| {
this.failure = Some(message);
cx.notify();
})
.is_err()
{
return;
}
}
}
}
}
Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => {
*last_token_usage.lock() = token_usage;
}
Ok(LanguageModelCompletionEvent::Stop(reason)) => {
if !failure_complete {
finish_with_status(
CodegenStatus::Error(anyhow!(
"Inline assistant stopped before completing \
`{REWRITE_SECTION_TOOL_NAME}` or \
`{FAILURE_MESSAGE_TOOL_NAME}` ({reason:?})"
)),
cx,
);
} else if matches!(reason, StopReason::MaxTokens | StopReason::Refusal) {
finish_with_status(
CodegenStatus::Error(anyhow!(
"Inline assistant response was interrupted ({reason:?})"
)),
cx,
);
} else {
finish_with_status(CodegenStatus::Done, cx);
}
return;
}
Ok(LanguageModelCompletionEvent::ToolUseJsonParseError {
tool_name,
json_parse_error,
..
}) => {
finish_with_status(
CodegenStatus::Error(anyhow!(
"Invalid input for inline assistant tool \
`{tool_name}`: {json_parse_error}"
)),
cx,
);
return;
}
Ok(
LanguageModelCompletionEvent::Queued { .. }
| LanguageModelCompletionEvent::Started
| LanguageModelCompletionEvent::Text(_)
| LanguageModelCompletionEvent::Thinking { .. }
| LanguageModelCompletionEvent::RedactedThinking { .. }
| LanguageModelCompletionEvent::ReasoningDetails(_)
| LanguageModelCompletionEvent::Compaction(_),
) => {}
Err(error) => {
finish_with_status(CodegenStatus::Error(error.into()), cx);
return;
}
}
};
let move_last_token_usage = last_token_usage.clone();
let move_rewrite_state = rewrite_state.clone();
let remaining_text_stream = completion_events.filter_map(move |e| {
let process_tool_use = process_tool_use.clone();
let last_token_usage = move_last_token_usage.clone();
let rewrite_state = move_rewrite_state.clone();
async move {
match e {
Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
let output = match process_tool_use(tool_use) {
Ok(Some(output)) => output,
Ok(None) => return None,
Err(error) => return Some(Err(error.into())),
};
match output {
ToolUseOutput::Rewrite { text } => Some(Ok(text)),
ToolUseOutput::Failure { .. } => Some(Err(anyhow!(
"Inline assistant called `{FAILURE_MESSAGE_TOOL_NAME}` \
after starting `{REWRITE_SECTION_TOOL_NAME}`"
)
.into())),
}
}
Ok(LanguageModelCompletionEvent::UsageUpdate(token_usage)) => {
*last_token_usage.lock() = token_usage;
None
}
Ok(LanguageModelCompletionEvent::ToolUseJsonParseError {
tool_name,
json_parse_error,
..
}) => Some(Err(anyhow!(
"Invalid input for inline assistant tool \
`{tool_name}`: {json_parse_error}"
)
.into())),
Ok(LanguageModelCompletionEvent::Stop(reason)) => {
let rewrite_complete = rewrite_state
.lock()
.as_ref()
.is_some_and(|state| state.is_input_complete);
if !rewrite_complete {
Some(Err(anyhow!(
"Inline assistant stopped before completing \
`{REWRITE_SECTION_TOOL_NAME}` ({reason:?})"
)
.into()))
} else if matches!(reason, StopReason::MaxTokens | StopReason::Refusal)
{
Some(Err(anyhow!(
"Inline assistant response was interrupted ({reason:?})"
)
.into()))
} else {
None
}
}
Ok(
LanguageModelCompletionEvent::Queued { .. }
| LanguageModelCompletionEvent::Started
| LanguageModelCompletionEvent::Text(_)
| LanguageModelCompletionEvent::Thinking { .. }
| LanguageModelCompletionEvent::RedactedThinking { .. }
| LanguageModelCompletionEvent::StartMessage { .. }
| LanguageModelCompletionEvent::ReasoningDetails(_)
| LanguageModelCompletionEvent::Compaction(_),
) => None,
Err(error) => Some(Err(error)),
}
}
});
let completion_status = futures::stream::once(async move {
if rewrite_state
.lock()
.as_ref()
.is_some_and(|state| state.is_input_complete)
{
None
} else {
Some(Err(anyhow!(
"Inline assistant response ended before completing \
`{REWRITE_SECTION_TOOL_NAME}`"
)
.into()))
}
})
.filter_map(futures::future::ready);
let text_stream = Box::pin(
futures::stream::once(async { Ok(first_text) })
.chain(remaining_text_stream)
.chain(completion_status),
);
let language_model_text_stream = LanguageModelTextStream {
message_id: message_id,
stream: text_stream,
last_token_usage,
};
let Some(task) = codegen
.update(cx, move |codegen, cx| {
codegen.handle_stream(
model,
/* strip_invalid_spans: */ false,
async { Ok(language_model_text_stream) },
cx,
)
})
.ok()
else {
return;
};
task.await;
})
}
}
#[derive(Copy, Clone, Debug)]
pub enum CodegenEvent {
Finished,
Undone,
}
struct StripInvalidSpans<T> {
stream: T,
stream_done: bool,
buffer: String,
first_line: bool,
line_end: bool,
starts_with_code_block: bool,
}
impl<T> StripInvalidSpans<T>
where
T: Stream<Item = Result<String>>,
{
fn new(stream: T) -> Self {
Self {
stream,
stream_done: false,
buffer: String::new(),
first_line: true,
line_end: false,
starts_with_code_block: false,
}
}
}
impl<T> Stream for StripInvalidSpans<T>
where
T: Stream<Item = Result<String>>,
{
type Item = Result<String>;
fn poll_next(self: Pin<&mut Self>, cx: &mut task::Context) -> Poll<Option<Self::Item>> {
const CODE_BLOCK_DELIMITER: &str = "```";
const CURSOR_SPAN: &str = "<|CURSOR|>";
let this = unsafe { self.get_unchecked_mut() };
loop {
if !this.stream_done {
let mut stream = unsafe { Pin::new_unchecked(&mut this.stream) };
match stream.as_mut().poll_next(cx) {
Poll::Ready(Some(Ok(chunk))) => {
this.buffer.push_str(&chunk);
}
Poll::Ready(Some(Err(error))) => return Poll::Ready(Some(Err(error))),
Poll::Ready(None) => {
this.stream_done = true;
}
Poll::Pending => return Poll::Pending,
}
}
let mut chunk = String::new();
let mut consumed = 0;
if !this.buffer.is_empty() {
let mut lines = this.buffer.split('\n').enumerate().peekable();
while let Some((line_ix, line)) = lines.next() {
if line_ix > 0 {
this.first_line = false;
}
if this.first_line {
let trimmed_line = line.trim();
if lines.peek().is_some() {
if trimmed_line.starts_with(CODE_BLOCK_DELIMITER) {
consumed += line.len() + 1;
this.starts_with_code_block = true;
continue;
}
} else if trimmed_line.is_empty()
|| prefixes(CODE_BLOCK_DELIMITER)
.any(|prefix| trimmed_line.starts_with(prefix))
{
break;
}
}
let line_without_cursor = line.replace(CURSOR_SPAN, "");
if lines.peek().is_some() {
if this.line_end {
chunk.push('\n');
}
chunk.push_str(&line_without_cursor);
this.line_end = true;
consumed += line.len() + 1;
} else if this.stream_done {
if !this.starts_with_code_block
|| !line_without_cursor.trim().ends_with(CODE_BLOCK_DELIMITER)
{
if this.line_end {
chunk.push('\n');
}
chunk.push_str(line);
}
consumed += line.len();
} else {
let trimmed_line = line.trim();
if trimmed_line.is_empty()
|| prefixes(CURSOR_SPAN).any(|prefix| trimmed_line.ends_with(prefix))
|| prefixes(CODE_BLOCK_DELIMITER)
.any(|prefix| trimmed_line.ends_with(prefix))
{
break;
} else {
if this.line_end {
chunk.push('\n');
this.line_end = false;
}
chunk.push_str(&line_without_cursor);
consumed += line.len();
}
}
}
}
this.buffer = this.buffer.split_off(consumed);
if !chunk.is_empty() {
return Poll::Ready(Some(Ok(chunk)));
} else if this.stream_done {
return Poll::Ready(None);
}
}
}
}
fn prefixes(text: &str) -> impl Iterator<Item = &str> {
(0..text.len() - 1).map(|ix| &text[..ix + 1])
}
#[derive(Default)]
pub struct Diff {
pub deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)>,
pub inserted_row_ranges: Vec<Range<Anchor>>,
}
impl Diff {
fn is_empty(&self) -> bool {
self.deleted_row_ranges.is_empty() && self.inserted_row_ranges.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures::{
Stream,
stream::{self},
};
use gpui::TestAppContext;
use indoc::indoc;
use language::{Buffer, Point};
use language_model::fake_provider::FakeLanguageModel;
use language_model::{
LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelRegistry,
LanguageModelToolUse, StopReason, TokenUsage,
};
use languages::rust_lang;
use rand::prelude::*;
use settings::SettingsStore;
use std::{future, sync::Arc};
#[gpui::test(iterations = 10)]
async fn test_transform_autoindent(cx: &mut TestAppContext, mut rng: StdRng) {
init_test(cx);
let text = indoc! {"
fn main() {
let x = 0;
for _ in 0..10 {
x += 1;
}
}
"};
let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
let range = buffer.read_with(cx, |buffer, cx| {
let snapshot = buffer.snapshot(cx);
snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(4, 5))
});
let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
let codegen = cx.new(|cx| {
CodegenAlternative::new(
buffer.clone(),
range.clone(),
true,
prompt_builder,
Uuid::new_v4(),
cx,
)
});
let chunks_tx = simulate_response_stream(&codegen, cx);
let mut new_text = concat!(
" let mut x = 0;\n",
" while x < 10 {\n",
" x += 1;\n",
" }",
);
while !new_text.is_empty() {
let max_len = cmp::min(new_text.len(), 10);
let len = rng.random_range(1..=max_len);
let (chunk, suffix) = new_text.split_at(len);
chunks_tx.unbounded_send(chunk.to_string()).unwrap();
new_text = suffix;
cx.background_executor.run_until_parked();
}
drop(chunks_tx);
cx.background_executor.run_until_parked();
assert_eq!(
buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
indoc! {"
fn main() {
let mut x = 0;
while x < 10 {
x += 1;
}
}
"}
);
}
#[gpui::test(iterations = 10)]
async fn test_autoindent_when_generating_past_indentation(
cx: &mut TestAppContext,
mut rng: StdRng,
) {
init_test(cx);
let text = indoc! {"
fn main() {
le
}
"};
let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
let range = buffer.read_with(cx, |buffer, cx| {
let snapshot = buffer.snapshot(cx);
snapshot.anchor_before(Point::new(1, 6))..snapshot.anchor_after(Point::new(1, 6))
});
let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
let codegen = cx.new(|cx| {
CodegenAlternative::new(
buffer.clone(),
range.clone(),
true,
prompt_builder,
Uuid::new_v4(),
cx,
)
});
let chunks_tx = simulate_response_stream(&codegen, cx);
cx.background_executor.run_until_parked();
let mut new_text = concat!(
"t mut x = 0;\n",
"while x < 10 {\n",
" x += 1;\n",
"}", //
);
while !new_text.is_empty() {
let max_len = cmp::min(new_text.len(), 10);
let len = rng.random_range(1..=max_len);
let (chunk, suffix) = new_text.split_at(len);
chunks_tx.unbounded_send(chunk.to_string()).unwrap();
new_text = suffix;
cx.background_executor.run_until_parked();
}
drop(chunks_tx);
cx.background_executor.run_until_parked();
assert_eq!(
buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
indoc! {"
fn main() {
let mut x = 0;
while x < 10 {
x += 1;
}
}
"}
);
}
#[gpui::test(iterations = 10)]
async fn test_autoindent_when_generating_before_indentation(
cx: &mut TestAppContext,
mut rng: StdRng,
) {
init_test(cx);
let text = concat!(
"fn main() {\n",
" \n",
"}\n" //
);
let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
let range = buffer.read_with(cx, |buffer, cx| {
let snapshot = buffer.snapshot(cx);
snapshot.anchor_before(Point::new(1, 2))..snapshot.anchor_after(Point::new(1, 2))
});
let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
let codegen = cx.new(|cx| {
CodegenAlternative::new(
buffer.clone(),
range.clone(),
true,
prompt_builder,
Uuid::new_v4(),
cx,
)
});
let chunks_tx = simulate_response_stream(&codegen, cx);
cx.background_executor.run_until_parked();
let mut new_text = concat!(
"let mut x = 0;\n",
"while x < 10 {\n",
" x += 1;\n",
"}", //
);
while !new_text.is_empty() {
let max_len = cmp::min(new_text.len(), 10);
let len = rng.random_range(1..=max_len);
let (chunk, suffix) = new_text.split_at(len);
chunks_tx.unbounded_send(chunk.to_string()).unwrap();
new_text = suffix;
cx.background_executor.run_until_parked();
}
drop(chunks_tx);
cx.background_executor.run_until_parked();
assert_eq!(
buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
indoc! {"
fn main() {
let mut x = 0;
while x < 10 {
x += 1;
}
}
"}
);
}
#[gpui::test(iterations = 10)]
async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
init_test(cx);
let text = indoc! {"
func main() {
\tx := 0
\tfor i := 0; i < 10; i++ {
\t\tx++
\t}
}
"};
let buffer = cx.new(|cx| Buffer::local(text, cx));
let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
let range = buffer.read_with(cx, |buffer, cx| {
let snapshot = buffer.snapshot(cx);
snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(4, 2))
});
let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
let codegen = cx.new(|cx| {
CodegenAlternative::new(
buffer.clone(),
range.clone(),
true,
prompt_builder,
Uuid::new_v4(),
cx,
)
});
let chunks_tx = simulate_response_stream(&codegen, cx);
let new_text = concat!(
"func main() {\n",
"\tx := 0\n",
"\tfor x < 10 {\n",
"\t\tx++\n",
"\t}", //
);
chunks_tx.unbounded_send(new_text.to_string()).unwrap();
drop(chunks_tx);
cx.background_executor.run_until_parked();
assert_eq!(
buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
indoc! {"
func main() {
\tx := 0
\tfor x < 10 {
\t\tx++
\t}
}
"}
);
}
#[gpui::test]
async fn test_inactive_codegen_alternative(cx: &mut TestAppContext) {
init_test(cx);
let text = indoc! {"
fn main() {
let x = 0;
}
"};
let buffer = cx.new(|cx| Buffer::local(text, cx).with_language(rust_lang(), cx));
let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
let range = buffer.read_with(cx, |buffer, cx| {
let snapshot = buffer.snapshot(cx);
snapshot.anchor_before(Point::new(1, 0))..snapshot.anchor_after(Point::new(1, 14))
});
let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
let codegen = cx.new(|cx| {
CodegenAlternative::new(
buffer.clone(),
range.clone(),
false,
prompt_builder,
Uuid::new_v4(),
cx,
)
});
let chunks_tx = simulate_response_stream(&codegen, cx);
chunks_tx
.unbounded_send("let mut x = 0;\nx += 1;".to_string())
.unwrap();
drop(chunks_tx);
cx.run_until_parked();
// The codegen is inactive, so the buffer doesn't get modified.
assert_eq!(
buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
text
);
// Activating the codegen applies the changes.
codegen.update(cx, |codegen, cx| codegen.set_active(true, cx));
assert_eq!(
buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
indoc! {"
fn main() {
let mut x = 0;
x += 1;
}
"}
);
// Deactivating the codegen undoes the changes.
codegen.update(cx, |codegen, cx| codegen.set_active(false, cx));
cx.run_until_parked();
assert_eq!(
buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
text
);
}
// When not streaming tool calls, we strip backticks as part of parsing the model's
// plain text response. This is a regression test for a bug where we stripped
// backticks incorrectly.
#[gpui::test]
async fn test_allows_model_to_output_backticks(cx: &mut TestAppContext) {
init_test(cx);
let text = "- Improved; `cmd+click` behavior. Now requires `cmd` to be pressed before the click starts or it doesn't run. ([#44579](https://github.com/zed-industries/zed/pull/44579); thanks [Zachiah](https://github.com/Zachiah))";
let buffer = cx.new(|cx| Buffer::local("", cx));
let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
let range = buffer.read_with(cx, |buffer, cx| {
let snapshot = buffer.snapshot(cx);
snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(0, 0))
});
let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
let codegen = cx.new(|cx| {
CodegenAlternative::new(
buffer.clone(),
range.clone(),
true,
prompt_builder,
Uuid::new_v4(),
cx,
)
});
let events_tx = simulate_tool_based_completion(&codegen, cx);
let chunk_len = text.find('`').unwrap();
events_tx
.unbounded_send(rewrite_tool_use("tool_1", &text[..chunk_len], false))
.unwrap();
events_tx
.unbounded_send(rewrite_tool_use("tool_1", &text, true))
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))
.unwrap();
drop(events_tx);
cx.run_until_parked();
assert_eq!(
buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
text
);
}
// Regression test: a second rewrite tool use with a *shorter* replacement_text
// than the first would cause an index-out-of-bounds panic because the
// chars_read_so_far counter was shared across all tool use IDs.
#[gpui::test]
async fn test_tool_completion_rejects_multiple_rewrite_tool_calls(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(rewrite_tool_use("tool_1", "longer replacement text", true))
.unwrap();
events_tx
.unbounded_send(rewrite_tool_use("tool_2", "short", true))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some("Inline assistant called `rewrite_section` multiple times")
);
}
#[gpui::test]
async fn test_tool_completion_ignores_non_tool_events(cx: &mut TestAppContext) {
let (buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(LanguageModelCompletionEvent::Queued { position: 1 })
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::Started)
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::StartMessage {
message_id: "message_1".into(),
})
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::ReasoningDetails(
serde_json::json!({ "encrypted_content": "reasoning" }),
))
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::Text(
"I'll update the section.".into(),
))
.unwrap();
events_tx
.unbounded_send(rewrite_tool_use("tool_1", "replace", false))
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::Thinking {
text: "Finishing the rewrite".into(),
signature: None,
})
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::RedactedThinking {
data: "redacted".into(),
})
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::UsageUpdate(TokenUsage {
output_tokens: 1,
..TokenUsage::default()
}))
.unwrap();
events_tx
.unbounded_send(rewrite_tool_use("tool_1", "replacement", true))
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))
.unwrap();
drop(events_tx);
cx.run_until_parked();
assert_eq!(
buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
"replacement"
);
assert!(codegen.read_with(cx, |codegen, _cx| matches!(
&codegen.status,
CodegenStatus::Done
)));
}
#[gpui::test]
async fn test_tool_completion_errors_when_stopped_without_tool_call(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(LanguageModelCompletionEvent::ReasoningDetails(
serde_json::json!([]),
))
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some(
"Inline assistant stopped before completing `rewrite_section` or \
`failure_message` (EndTurn)"
)
);
}
#[gpui::test]
async fn test_tool_completion_errors_on_eof_without_tool_call(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some(
"Inline assistant response ended without completing `rewrite_section` or \
`failure_message`"
)
);
}
#[gpui::test]
async fn test_tool_completion_accepts_failure_message(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(failure_tool_use("The selection cannot be rewritten.", true))
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))
.unwrap();
drop(events_tx);
cx.run_until_parked();
assert!(codegen.read_with(cx, |codegen, _cx| matches!(
&codegen.status,
CodegenStatus::Done
)));
assert_eq!(
codegen.read_with(cx, |codegen, _cx| codegen.current_failure()),
Some("The selection cannot be rewritten.".into())
);
}
#[gpui::test]
async fn test_tool_completion_errors_on_invalid_tool_json(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(LanguageModelCompletionEvent::ToolUseJsonParseError {
id: "tool_1".into(),
tool_name: REWRITE_SECTION_TOOL_NAME.into(),
raw_input: r#"{"replacement_text":"#.into(),
json_parse_error: "EOF while parsing a string".into(),
})
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some(
"Invalid input for inline assistant tool `rewrite_section`: \
EOF while parsing a string"
)
);
}
#[gpui::test]
async fn test_tool_completion_errors_when_incomplete_rewrite_stops(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(rewrite_tool_use("tool_1", "partial", false))
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::MaxTokens))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some(
"Inline assistant stopped before completing `rewrite_section` \
(MaxTokens)"
)
);
}
#[gpui::test]
async fn test_tool_completion_errors_when_incomplete_rewrite_reaches_eof(
cx: &mut TestAppContext,
) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(rewrite_tool_use("tool_1", "partial", false))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some("Inline assistant response ended before completing `rewrite_section`")
);
}
#[gpui::test]
async fn test_tool_completion_errors_when_interrupted_after_complete_rewrite(
cx: &mut TestAppContext,
) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(rewrite_tool_use("tool_1", "replacement", true))
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::MaxTokens))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some("Inline assistant response was interrupted (MaxTokens)")
);
}
#[gpui::test]
async fn test_tool_completion_errors_when_interrupted_after_failure_message(
cx: &mut TestAppContext,
) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(failure_tool_use("The selection cannot be rewritten.", true))
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::MaxTokens))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some("Inline assistant response was interrupted (MaxTokens)")
);
}
#[gpui::test]
async fn test_tool_completion_requires_failure_message_to_complete(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(failure_tool_use("Partial failure", false))
.unwrap();
events_tx
.unbounded_send(LanguageModelCompletionEvent::Stop(StopReason::EndTurn))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some(
"Inline assistant stopped before completing `rewrite_section` or \
`failure_message` (EndTurn)"
)
);
assert_eq!(
codegen.read_with(cx, |codegen, _cx| codegen.current_failure()),
None
);
}
#[gpui::test]
async fn test_tool_completion_rejects_failure_after_rewrite(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(rewrite_tool_use("tool_1", "replacement", true))
.unwrap();
events_tx
.unbounded_send(failure_tool_use("Cannot rewrite selection", true))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some(
"Inline assistant called `failure_message` after starting \
`rewrite_section`"
)
);
}
#[gpui::test]
async fn test_tool_completion_rejects_rewrite_after_failure(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(failure_tool_use("Cannot rewrite selection", true))
.unwrap();
events_tx
.unbounded_send(rewrite_tool_use("tool_1", "replacement", true))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some("Inline assistant called `rewrite_section` after `failure_message`")
);
}
#[gpui::test]
async fn test_tool_completion_rejects_rewrite_after_partial_failure(cx: &mut TestAppContext) {
let (buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(failure_tool_use("Cannot rewrite", false))
.unwrap();
events_tx
.unbounded_send(rewrite_tool_use("tool_1", "replacement", true))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(
error.as_deref(),
Some("Inline assistant called `rewrite_section` after `failure_message`")
);
assert_eq!(
buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
""
);
}
#[gpui::test]
async fn test_tool_completion_propagates_stream_error_before_rewrite(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion_results(&codegen, cx);
events_tx
.unbounded_send(Err(anyhow!("provider stream failed").into()))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(error.as_deref(), Some("provider stream failed"));
}
#[gpui::test]
async fn test_tool_completion_propagates_stream_error_after_rewrite(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion_results(&codegen, cx);
events_tx
.unbounded_send(Ok(rewrite_tool_use("tool_1", "replacement", true)))
.unwrap();
events_tx
.unbounded_send(Err(anyhow!("provider stream failed").into()))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert_eq!(error.as_deref(), Some("provider stream failed"));
}
#[gpui::test]
async fn test_tool_completion_errors_on_invalid_tool_input(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(LanguageModelCompletionEvent::ToolUse(
LanguageModelToolUse {
id: "tool_1".into(),
name: REWRITE_SECTION_TOOL_NAME.into(),
raw_input: r#"{"replacement_text":42}"#.into(),
input: language_model::LanguageModelToolUseInput::Json(
serde_json::json!({ "replacement_text": 42 }),
),
is_input_complete: true,
thought_signature: None,
},
))
.unwrap();
drop(events_tx);
cx.run_until_parked();
let error = codegen.read_with(cx, |codegen, _cx| match &codegen.status {
CodegenStatus::Error(error) => Some(error.to_string()),
_ => None,
});
assert!(
error
.as_deref()
.is_some_and(|error| error.starts_with("invalid input for `rewrite_section` tool"))
);
}
#[gpui::test]
async fn test_start_clears_previous_failure(cx: &mut TestAppContext) {
let (_buffer, codegen) = new_tool_based_codegen(cx);
cx.update(|cx| AgentSettings::register(cx));
let events_tx = simulate_tool_based_completion(&codegen, cx);
events_tx
.unbounded_send(failure_tool_use("Cannot rewrite selection", true))
.unwrap();
drop(events_tx);
cx.run_until_parked();
assert_eq!(
codegen.read_with(cx, |codegen, _cx| codegen.current_failure()),
Some("Cannot rewrite selection".into())
);
codegen.update(cx, |codegen, cx| {
codegen
.start(
"Rewrite the selection".into(),
Task::ready(None).shared(),
Arc::new(FakeLanguageModel::default()),
cx,
)
.unwrap();
});
assert_eq!(
codegen.read_with(cx, |codegen, _cx| codegen.current_failure()),
None
);
}
#[gpui::test]
async fn test_strip_invalid_spans_from_codeblock() {
assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;
assert_chunks("```\nLorem ipsum dolor", "Lorem ipsum dolor").await;
assert_chunks("```\nLorem ipsum dolor\n```", "Lorem ipsum dolor").await;
assert_chunks(
"```html\n```js\nLorem ipsum dolor\n```\n```",
"```js\nLorem ipsum dolor\n```",
)
.await;
assert_chunks("``\nLorem ipsum dolor\n```", "``\nLorem ipsum dolor\n```").await;
assert_chunks("Lorem<|CURSOR|> ipsum", "Lorem ipsum").await;
assert_chunks("Lorem ipsum", "Lorem ipsum").await;
assert_chunks("```\n<|CURSOR|>Lorem ipsum\n```", "Lorem ipsum").await;
async fn assert_chunks(text: &str, expected_text: &str) {
for chunk_size in 1..=text.len() {
let actual_text = StripInvalidSpans::new(chunks(text, chunk_size))
.map(|chunk| chunk.unwrap())
.collect::<String>()
.await;
assert_eq!(
actual_text, expected_text,
"failed to strip invalid spans, chunk size: {}",
chunk_size
);
}
}
fn chunks(text: &str, size: usize) -> impl Stream<Item = Result<String>> {
stream::iter(
text.chars()
.collect::<Vec<_>>()
.chunks(size)
.map(|chunk| Ok(chunk.iter().collect::<String>()))
.collect::<Vec<_>>(),
)
}
}
fn init_test(cx: &mut TestAppContext) {
cx.update(LanguageModelRegistry::test);
cx.set_global(cx.update(SettingsStore::test));
}
fn new_tool_based_codegen(
cx: &mut TestAppContext,
) -> (Entity<MultiBuffer>, Entity<CodegenAlternative>) {
init_test(cx);
let buffer = cx.new(|cx| Buffer::local("", cx));
let buffer = cx.new(|cx| MultiBuffer::singleton(buffer, cx));
let range = buffer.read_with(cx, |buffer, cx| {
let snapshot = buffer.snapshot(cx);
snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(0, 0))
});
let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());
let codegen = cx.new(|cx| {
CodegenAlternative::new(
buffer.clone(),
range,
true,
prompt_builder,
Uuid::new_v4(),
cx,
)
});
(buffer, codegen)
}
fn simulate_response_stream(
codegen: &Entity<CodegenAlternative>,
cx: &mut TestAppContext,
) -> mpsc::UnboundedSender<String> {
let (chunks_tx, chunks_rx) = mpsc::unbounded();
let model = Arc::new(FakeLanguageModel::default());
codegen.update(cx, |codegen, cx| {
codegen.generation = codegen.handle_stream(
model,
/* strip_invalid_spans: */ false,
future::ready(Ok(LanguageModelTextStream {
message_id: None,
stream: chunks_rx.map(Ok).boxed(),
last_token_usage: Arc::new(Mutex::new(TokenUsage::default())),
})),
cx,
);
});
chunks_tx
}
fn simulate_tool_based_completion_results(
codegen: &Entity<CodegenAlternative>,
cx: &mut TestAppContext,
) -> mpsc::UnboundedSender<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
{
let (events_tx, events_rx) = mpsc::unbounded();
let model = Arc::new(FakeLanguageModel::default());
codegen.update(cx, |codegen, cx| {
codegen.generation =
codegen.handle_completion(model, Task::ready(Ok(events_rx.boxed())), cx);
});
events_tx
}
fn simulate_tool_based_completion(
codegen: &Entity<CodegenAlternative>,
cx: &mut TestAppContext,
) -> mpsc::UnboundedSender<LanguageModelCompletionEvent> {
let (events_tx, events_rx) = mpsc::unbounded();
let model = Arc::new(FakeLanguageModel::default());
codegen.update(cx, |codegen, cx| {
let completion_stream = Task::ready(Ok(events_rx.map(Ok).boxed()
as BoxStream<
'static,
Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
>));
codegen.generation = codegen.handle_completion(model, completion_stream, cx);
});
events_tx
}
fn failure_tool_use(message: &str, is_input_complete: bool) -> LanguageModelCompletionEvent {
let input = FailureMessageInput {
message: message.into(),
};
LanguageModelCompletionEvent::ToolUse(LanguageModelToolUse {
id: "failure_tool".into(),
name: FAILURE_MESSAGE_TOOL_NAME.into(),
raw_input: serde_json::to_string(&input).unwrap(),
input: language_model::LanguageModelToolUseInput::Json(
serde_json::to_value(&input).unwrap(),
),
is_input_complete,
thought_signature: None,
})
}
fn rewrite_tool_use(
id: &str,
replacement_text: &str,
is_complete: bool,
) -> LanguageModelCompletionEvent {
let input = RewriteSectionInput {
replacement_text: replacement_text.into(),
};
LanguageModelCompletionEvent::ToolUse(LanguageModelToolUse {
id: id.into(),
name: REWRITE_SECTION_TOOL_NAME.into(),
raw_input: serde_json::to_string(&input).unwrap(),
input: language_model::LanguageModelToolUseInput::Json(
serde_json::to_value(&input).unwrap(),
),
is_input_complete: is_complete,
thought_signature: None,
})
}
}
Reference in new issue View Git Blame Copy Permalink