Skip to content

Commit 711c5af

Browse files
kennvene-awskennvene-aws
committed
feat: Adds continuation id logic (#3114)
* feat: New continuation id logic * chore: Centralize metadata * chore: Moves struct to right place --------- Co-authored-by: Kenneth S. <[email protected]>
1 parent 87575bd commit 711c5af

File tree

7 files changed

+237
-73
lines changed

7 files changed

+237
-73
lines changed

‎crates/chat-cli/src/api_client/mod.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -382,6 +382,7 @@ impl ApiClient {
382382
conversation_id,
383383
user_input_message,
384384
history,
385+
agent_continuation_id,
385386
} = conversation;
386387

387388
let model_id_opt: Option<String> = user_input_message.model_id.clone();
@@ -400,6 +401,8 @@ impl ApiClient {
400401
.map(|v| v.into_iter().map(|i| i.try_into()).collect::<Result<Vec<_>, _>>())
401402
.transpose()?,
402403
)
404+
.set_agent_continuation_id(agent_continuation_id)
405+
.agent_task_type(amzn_codewhisperer_streaming_client::types::AgentTaskType::Vibe)
403406
.build()
404407
.expect("building conversation should not fail");
405408

@@ -744,6 +747,7 @@ mod tests {
744747
model_id: Some("model".to_owned()),
745748
},
746749
history: None,
750+
agent_continuation_id: None,
747751
})
748752
.await
749753
.unwrap();

‎crates/chat-cli/src/api_client/model.rs‎

Lines changed: 42 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ pub struct ConversationState {
9797
pub conversation_id: Option<String>,
9898
pub user_input_message: UserInputMessage,
9999
pub history: Option<Vec<ChatMessage>>,
100+
pub agent_continuation_id: Option<String>,
100101
}
101102

102103
#[derive(Debug, Clone)]
@@ -542,7 +543,7 @@ impl TryFrom<AssistantResponseMessage> for amzn_qdeveloper_streaming_client::typ
542543
}
543544

544545
#[non_exhaustive]
545-
#[derive(Debug, Clone, PartialEq, Eq)]
546+
#[derive(Debug, Clone, PartialEq)]
546547
pub enum ChatResponseStream {
547548
AssistantResponseEvent {
548549
content: String,
@@ -563,6 +564,16 @@ pub enum ChatResponseStream {
563564
conversation_id: Option<String>,
564565
utterance_id: Option<String>,
565566
},
567+
MetadataEvent {
568+
total_tokens: Option<i32>,
569+
uncached_input_tokens: Option<i32>,
570+
output_tokens: Option<i32>,
571+
cache_read_input_tokens: Option<i32>,
572+
cache_write_input_tokens: Option<i32>,
573+
},
574+
MeteringEvent {
575+
usage: Option<f64>,
576+
},
566577
SupplementaryWebLinksEvent(()),
567578
ToolUseEvent {
568579
tool_use_id: String,
@@ -609,6 +620,18 @@ impl From<amzn_codewhisperer_streaming_client::types::ChatResponseStream> for Ch
609620
conversation_id,
610621
utterance_id,
611622
},
623+
amzn_codewhisperer_streaming_client::types::ChatResponseStream::MetadataEvent(
624+
amzn_codewhisperer_streaming_client::types::MetadataEvent { token_usage, .. },
625+
) => ChatResponseStream::MetadataEvent {
626+
total_tokens: token_usage.as_ref().map(|t| t.total_tokens),
627+
uncached_input_tokens: token_usage.as_ref().map(|t| t.uncached_input_tokens),
628+
output_tokens: token_usage.as_ref().map(|t| t.output_tokens),
629+
cache_read_input_tokens: token_usage.as_ref().and_then(|t| t.cache_read_input_tokens),
630+
cache_write_input_tokens: token_usage.as_ref().and_then(|t| t.cache_write_input_tokens),
631+
},
632+
amzn_codewhisperer_streaming_client::types::ChatResponseStream::MeteringEvent(
633+
amzn_codewhisperer_streaming_client::types::MeteringEvent { usage, .. },
634+
) => ChatResponseStream::MeteringEvent { usage },
612635
amzn_codewhisperer_streaming_client::types::ChatResponseStream::ToolUseEvent(
613636
amzn_codewhisperer_streaming_client::types::ToolUseEvent {
614637
tool_use_id,
@@ -626,7 +649,7 @@ impl From<amzn_codewhisperer_streaming_client::types::ChatResponseStream> for Ch
626649
amzn_codewhisperer_streaming_client::types::ChatResponseStream::SupplementaryWebLinksEvent(_) => {
627650
ChatResponseStream::SupplementaryWebLinksEvent(())
628651
},
629-
_ => ChatResponseStream::Unknown,
652+
_other => ChatResponseStream::Unknown,
630653
}
631654
}
632655
}
@@ -665,6 +688,18 @@ impl From<amzn_qdeveloper_streaming_client::types::ChatResponseStream> for ChatR
665688
conversation_id,
666689
utterance_id,
667690
},
691+
amzn_qdeveloper_streaming_client::types::ChatResponseStream::MetadataEvent(
692+
amzn_qdeveloper_streaming_client::types::MetadataEvent { token_usage, .. },
693+
) => ChatResponseStream::MetadataEvent {
694+
total_tokens: token_usage.as_ref().map(|t| t.total_tokens),
695+
uncached_input_tokens: token_usage.as_ref().map(|t| t.uncached_input_tokens),
696+
output_tokens: token_usage.as_ref().map(|t| t.output_tokens),
697+
cache_read_input_tokens: token_usage.as_ref().and_then(|t| t.cache_read_input_tokens),
698+
cache_write_input_tokens: token_usage.as_ref().and_then(|t| t.cache_write_input_tokens),
699+
},
700+
amzn_qdeveloper_streaming_client::types::ChatResponseStream::MeteringEvent(
701+
amzn_qdeveloper_streaming_client::types::MeteringEvent { usage, .. },
702+
) => ChatResponseStream::MeteringEvent { usage },
668703
amzn_qdeveloper_streaming_client::types::ChatResponseStream::ToolUseEvent(
669704
amzn_qdeveloper_streaming_client::types::ToolUseEvent {
670705
tool_use_id,
@@ -682,7 +717,7 @@ impl From<amzn_qdeveloper_streaming_client::types::ChatResponseStream> for ChatR
682717
amzn_qdeveloper_streaming_client::types::ChatResponseStream::SupplementaryWebLinksEvent(_) => {
683718
ChatResponseStream::SupplementaryWebLinksEvent(())
684719
},
685-
_ => ChatResponseStream::Unknown,
720+
_other => ChatResponseStream::Unknown,
686721
}
687722
}
688723
}
@@ -870,7 +905,8 @@ impl From<UserInputMessage> for amzn_codewhisperer_streaming_client::types::User
870905
.set_user_input_message_context(value.user_input_message_context.map(Into::into))
871906
.set_user_intent(value.user_intent.map(Into::into))
872907
.set_model_id(value.model_id)
873-
.origin(amzn_codewhisperer_streaming_client::types::Origin::Cli)
908+
//TODO: Setup new origin.
909+
.origin(amzn_codewhisperer_streaming_client::types::Origin::AiEditor)
874910
.build()
875911
.expect("Failed to build UserInputMessage")
876912
}
@@ -884,7 +920,8 @@ impl From<UserInputMessage> for amzn_qdeveloper_streaming_client::types::UserInp
884920
.set_user_input_message_context(value.user_input_message_context.map(Into::into))
885921
.set_user_intent(value.user_intent.map(Into::into))
886922
.set_model_id(value.model_id)
887-
.origin(amzn_qdeveloper_streaming_client::types::Origin::Cli)
923+
//TODO: Setup new origin.
924+
.origin(amzn_qdeveloper_streaming_client::types::Origin::AiEditor)
888925
.build()
889926
.expect("Failed to build UserInputMessage")
890927
}

‎crates/chat-cli/src/cli/chat/context.rs‎

Lines changed: 12 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -378,14 +378,12 @@ async fn process_path(
378378
async fn add_file_to_context(os: &Os, path: &Path, context_files: &mut Vec<(String, String)>) -> Result<()> {
379379
let filename = path.to_string_lossy().to_string();
380380
let content = os.fs.read_to_string(path).await?;
381-
381+
382382
// Check if this is a steering file that needs front matter filtering
383-
if filename.contains(".kiro/steering") && filename.ends_with(".md") {
384-
if !should_include_steering_file(&content)? {
385-
return Ok(());
386-
}
383+
if filename.contains(".kiro/steering") && filename.ends_with(".md") && !should_include_steering_file(&content)? {
384+
return Ok(());
387385
}
388-
386+
389387
context_files.push((filename, content));
390388
Ok(())
391389
}
@@ -402,30 +400,30 @@ fn should_include_steering_file(content: &str) -> Result<bool> {
402400
// No front matter - include the file
403401
return Ok(true);
404402
}
405-
403+
406404
// Find the end of the front matter
407405
let lines: Vec<&str> = content.lines().collect();
408406
let mut end_index = None;
409-
407+
410408
for (i, line) in lines.iter().enumerate().skip(1) {
411409
if line.trim() == "---" {
412410
end_index = Some(i);
413411
break;
414412
}
415413
}
416-
414+
417415
let end_index = match end_index {
418416
Some(idx) => idx,
419417
None => {
420418
// Malformed front matter - include the file
421419
return Ok(true);
422-
}
420+
},
423421
};
424-
422+
425423
// Extract and parse the front matter
426424
let front_matter_lines = &lines[1..end_index];
427425
let front_matter_yaml = front_matter_lines.join("\n");
428-
426+
429427
match serde_yaml::from_str::<FrontMatter>(&front_matter_yaml) {
430428
Ok(front_matter) => {
431429
match front_matter.inclusion.as_deref() {
@@ -435,11 +433,11 @@ fn should_include_steering_file(content: &str) -> Result<bool> {
435433
None => Ok(true), // No inclusion field - include
436434
Some(_) => Ok(true), // Unknown inclusion value - include
437435
}
438-
}
436+
},
439437
Err(_) => {
440438
// Failed to parse front matter - include the file
441439
Ok(true)
442-
}
440+
},
443441
}
444442
}
445443

‎crates/chat-cli/src/cli/chat/conversation.rs‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,57 @@ pub struct McpServerInfo {
102102
pub config: CustomToolConfig,
103103
}
104104

105+
#[derive(Debug, Clone, Serialize, Deserialize)]
106+
pub struct UserTurnMetadata {
107+
continuation_id: String,
108+
/// [RequestMetadata] about the ongoing operation.
109+
requests: Vec<RequestMetadata>,
110+
}
111+
112+
impl Default for UserTurnMetadata {
113+
fn default() -> Self {
114+
Self::new()
115+
}
116+
}
117+
118+
/// Enum used to store metadata about user turns
119+
impl UserTurnMetadata {
120+
pub fn new() -> Self {
121+
Self {
122+
continuation_id: uuid::Uuid::new_v4().to_string(),
123+
requests: vec![],
124+
}
125+
}
126+
127+
pub fn continuation_id(&self) -> &str {
128+
&self.continuation_id
129+
}
130+
131+
pub fn add_request(&mut self, request: RequestMetadata) {
132+
self.requests.push(request);
133+
}
134+
135+
pub fn first_request(&self) -> Option<RequestMetadata> {
136+
self.requests.first().cloned()
137+
}
138+
139+
pub fn last_request(&self) -> Option<RequestMetadata> {
140+
self.requests.last().cloned()
141+
}
142+
143+
pub fn iter(&self) -> impl Iterator<Item = &RequestMetadata> {
144+
self.requests.iter()
145+
}
146+
147+
pub fn first(&self) -> Option<&RequestMetadata> {
148+
self.requests.first()
149+
}
150+
151+
pub fn last(&self) -> Option<&RequestMetadata> {
152+
self.requests.last()
153+
}
154+
}
155+
105156
/// Tracks state related to an ongoing conversation.
106157
#[derive(Debug, Clone, Serialize, Deserialize)]
107158
pub struct ConversationState {
@@ -149,6 +200,9 @@ pub struct ConversationState {
149200
/// Tangent mode checkpoint - stores main conversation when in tangent mode
150201
#[serde(default, skip_serializing_if = "Option::is_none")]
151202
tangent_state: Option<ConversationCheckpoint>,
203+
/// Metadata about the ongoing user turn operation
204+
#[serde(default)]
205+
pub user_turn_metadata: UserTurnMetadata,
152206
}
153207

154208
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -212,6 +266,7 @@ impl ConversationState {
212266
checkpoint_manager: None,
213267
mcp_enabled,
214268
tangent_state: None,
269+
user_turn_metadata: UserTurnMetadata::new(),
215270
}
216271
}
217272

@@ -378,6 +433,10 @@ impl ConversationState {
378433
self.next_message = None;
379434
}
380435

436+
pub fn current_continuation_id(&self) -> &str {
437+
self.user_turn_metadata.continuation_id()
438+
}
439+
381440
pub async fn set_next_user_message(&mut self, input: String) {
382441
debug_assert!(self.next_message.is_none(), "next_message should not exist");
383442
if let Some(next_message) = self.next_message.as_ref() {
@@ -614,6 +673,7 @@ impl ConversationState {
614673
dropped_context_files,
615674
tools: &self.tools,
616675
model_id: self.model_info.as_ref().map(|m| m.model_id.as_str()),
676+
continuation_id: Some(self.user_turn_metadata.continuation_id()),
617677
})
618678
}
619679

@@ -719,6 +779,7 @@ impl ConversationState {
719779
.unwrap_or(UserMessage::new_prompt(summary_content, None)) // should not happen
720780
.into_user_input_message(self.model_info.as_ref().map(|m| m.model_id.clone()), &tools),
721781
history: Some(flatten_history(history.iter())),
782+
agent_continuation_id: Some(self.user_turn_metadata.continuation_id().to_string()),
722783
})
723784
}
724785

@@ -777,6 +838,7 @@ Return only the JSON configuration, no additional text.",
777838
conversation_id: Some(self.conversation_id.clone()),
778839
user_input_message: generation_message.into_user_input_message(self.model.clone(), &tools),
779840
history: Some(flatten_history(history.iter())),
841+
agent_continuation_id: Some(self.user_turn_metadata.continuation_id().to_string()),
780842
})
781843
}
782844

@@ -992,6 +1054,7 @@ pub struct BackendConversationStateImpl<'a, T, U> {
9921054
pub dropped_context_files: Vec<(String, String)>,
9931055
pub tools: &'a HashMap<ToolOrigin, Vec<Tool>>,
9941056
pub model_id: Option<&'a str>,
1057+
pub continuation_id: Option<&'a str>,
9951058
}
9961059

9971060
impl BackendConversationStateImpl<'_, std::collections::vec_deque::Iter<'_, HistoryEntry>, Option<Vec<HistoryEntry>>> {
@@ -1007,6 +1070,7 @@ impl BackendConversationStateImpl<'_, std::collections::vec_deque::Iter<'_, Hist
10071070
conversation_id: Some(self.conversation_id.to_string()),
10081071
user_input_message,
10091072
history: Some(history),
1073+
agent_continuation_id: self.continuation_id.map(str::to_string),
10101074
})
10111075
}
10121076

0 commit comments

Comments
 (0)